t5de-proxy/internal/proxy/connect.go

package proxy

import (
	"bytes"
	"crypto/tls"
	"fmt"
	"io"
	"net"
	"net/http"
	"sync"
	"time"

	"git.t5de.dev/t5de/t5de-proxy/internal/imq"
)

// serveConnect connects the client's plaintext stream to host:port over TLS.
func (s *Server) serveConnect(w http.ResponseWriter, r *http.Request) {
	id := s.requests.Add(1)
	start := time.Now()
	target := r.Host
	host, _, err := net.SplitHostPort(target)
	if err != nil {
		http.Error(w, "t5de-proxy: expected CONNECT host:port", http.StatusBadRequest)
		return
	}
	config := &tls.Config{}
	if s.TLSConfig != nil {
		config = s.TLSConfig.Clone()
	}
	config.ServerName = host
	dial := s.Dial
	if dial == nil {
		dial = (&net.Dialer{Timeout: 30 * time.Second, KeepAlive: 30 * time.Second}).DialContext
	}
	raw, err := dial(r.Context(), "tcp", target)
	var upstream *tls.Conn
	if err == nil {
		upstream = tls.Client(raw, config)
		if err = upstream.HandshakeContext(r.Context()); err != nil {
			raw.Close()
		}
	}
	if err != nil {
		s.logf("#%d CONNECT %s failed after %s: %v", id, target, since(start), err)
		http.Error(w, "t5de-proxy: "+err.Error(), http.StatusBadGateway)
		return
	}
	defer upstream.Close()

	conn, buffered, err := http.NewResponseController(w).Hijack()
	if err != nil {
		s.logf("#%d CONNECT %s: %v", id, target, err)
		return
	}
	defer conn.Close()
	if _, err := io.WriteString(conn, "HTTP/1.1 200 Connection established\r\n\r\n"); err != nil {
		return
	}
	s.logf("#%d CONNECT %s open (%s)", id, target, since(start))

	// Whatever the client sent after its CONNECT and the server already read is the stream's start.
	var client io.Reader = conn
	if n := buffered.Reader.Buffered(); n > 0 {
		head, _ := buffered.Reader.Peek(n)
		client = io.MultiReader(bytes.NewReader(bytes.Clone(head)), conn)
	}

	// Frames go to the client whole, from the gateway or from the proxy answering a dropped message, so they
	// never interleave.
	toClient := &lockedWriter{w: conn}
	down := &stream{s: s, id: id, name: "g2c"}
	up := &stream{s: s, id: id, name: "c2g", filter: s.IMQFilter, reply: toClient}
	if w, ok := s.IMQFilter.(IMQWatcher); ok {
		c := &IMQConn{ID: id, w: toClient}
		down.watcher, down.conn = w, c
		up.watcher, up.conn = w, c
		defer w.Closed(c)
	}
	// IMQ is a session: when either side closes, it's over, so both connections close.
	var once sync.Once
	closeBoth := func() {
		once.Do(func() {
			conn.Close()
			upstream.Close()
		})
	}
	var wg sync.WaitGroup
	wg.Add(2)
	go func() {
		defer wg.Done()
		defer closeBoth()
		up.copy(upstream, client)
	}()
	go func() {
		defer wg.Done()
		defer closeBoth()
		down.copy(toClient, upstream)
	}()
	wg.Wait()
	s.logf("#%d CONNECT %s closed after %s: %s, %s", id, target, since(start), up.summary(), down.summary())
}

// IMQFilter sees the client's outgoing IMQ messages and can change or drop them.
type IMQFilter interface {
	// Outgoing returns the message to send instead of message, or drop to send nothing and tell the client it
	// was sent.
	Outgoing(queue, mount string, message []byte) (out []byte, drop bool)
}

// IMQWatcher is an IMQFilter that also follows each IMQ connection: the queues the client subscribes to, the
// messages the gateway delivers (which it can change), and it can deliver messages of its own.
type IMQWatcher interface {
	IMQFilter
	// Subscribed and Unsubscribed see the client's subscriptions on a connection.
	Subscribed(c *IMQConn, queues []string)
	Unsubscribed(c *IMQConn, queues []string)
	// Incoming returns the message to deliver instead of m's (m.Message itself to leave it).
	Incoming(c *IMQConn, m imq.Delivered) []byte
	// Closed: the connection is over (and with it, its subscriptions).
	Closed(c *IMQConn)
}

// IMQConn is one IMQ connection of the client's.
type IMQConn struct {
	ID int64
	w  io.Writer
}

// NewIMQConn is a connection writing frames to w (for tests of watchers).
func NewIMQConn(id int64, w io.Writer) *IMQConn { return &IMQConn{ID: id, w: w} }

// Deliver gives the client a message as if the gateway had delivered it.
func (c *IMQConn) Deliver(m imq.Delivered) error {
	_, err := c.w.Write(imq.EncodeFrame(imq.G2cSendMessage, m.Marshal()))
	return err
}

// lockedWriter serializes writes, each a whole number of frames.
type lockedWriter struct {
	mu sync.Mutex
	w  io.Writer
}

func (l *lockedWriter) Write(b []byte) (int, error) {
	l.mu.Lock()
	defer l.mu.Unlock()
	return l.w.Write(b)
}

// stream is one direction of a CONNECT. It passes IMQ frames on whole, counting them (and, outgoing, filtering
// them), or bytes as they come once the stream turns out not to be IMQ.
type stream struct {
	s       *Server
	id      int64
	name    string
	bytes   int64
	frames  int
	other   bool // not IMQ framing
	pending []byte
	filter  IMQFilter // outgoing only
	reply   io.Writer // outgoing only: to the client
	watcher IMQWatcher
	conn    *IMQConn
}

func (st *stream) copy(dst io.Writer, src io.Reader) {
	buf := make([]byte, 32<<10)
	for {
		n, err := src.Read(buf)
		if n > 0 {
			st.bytes += int64(n)
			if werr := st.pass(dst, buf[:n]); werr != nil {
				return
			}
		}
		if err != nil {
			return
		}
	}
}

// pass writes on the frames data completes.
func (st *stream) pass(dst io.Writer, data []byte) error {
	if st.other {
		_, err := dst.Write(data)
		return err
	}
	st.pending = append(st.pending, data...)
	var out []byte
	for {
		f, n, err := imq.Next(st.pending)
		if err != nil {
			st.other = true
			st.s.logf("#%d %s isn't IMQ framing; passing it through", st.id, st.name)
			out = append(out, st.pending...)
			st.pending = nil
			break
		}
		if n == 0 {
			break
		}
		frame := st.pending[:n]
		st.pending = st.pending[n:]
		st.frames++
		if st.s.Trace {
			st.s.logf("#%d %s %s %dB", st.id, st.name, imq.TypeName(f.Type), f.Size)
		}
		out = append(out, st.filtered(f, frame)...)
	}
	if len(st.pending) == 0 {
		st.pending = nil
	}
	if len(out) == 0 {
		return nil
	}
	_, err := dst.Write(out)
	return err
}

// filtered is the frame to send on: as is, changed by the filter, or none (answered here).
func (st *stream) filtered(f imq.Frame, frame []byte) []byte {
	if st.watcher != nil {
		switch f.Type {
		case imq.C2gSubscribe, imq.C2gUnsubscribe:
			if queues, err := imq.ParseSubscription(f.Data); err == nil {
				if f.Type == imq.C2gSubscribe {
					st.watcher.Subscribed(st.conn, queues)
				} else {
					st.watcher.Unsubscribed(st.conn, queues)
				}
			}
			return frame
		case imq.G2cSendMessage:
			m, err := imq.ParseDelivered(f.Data)
			if err != nil {
				return frame
			}
			out := st.watcher.Incoming(st.conn, m)
			if bytes.Equal(out, m.Message) {
				return frame
			}
			m.Message = out
			return imq.EncodeFrame(imq.G2cSendMessage, m.Marshal())
		}
	}
	if st.filter == nil || f.Type != imq.C2gSendMessage {
		return frame
	}
	m, err := imq.ParseSendMessage(f.Data)
	if err != nil {
		return frame
	}
	out, drop := st.filter.Outgoing(m.Queue, m.Mount, m.Message)
	if drop {
		if st.reply != nil {
			st.reply.Write(imq.EncodeFrame(imq.G2cResult, imq.Result(m.OpID, 0)))
		}
		return nil
	}
	if bytes.Equal(out, m.Message) {
		return frame
	}
	m.Message = out
	return imq.EncodeFrame(imq.C2gSendMessage, m.Marshal())
}

func (st *stream) summary() string {
	if st.other {
		return fmt.Sprintf("%s %dB", st.name, st.bytes)
	}
	return fmt.Sprintf("%s %dB in %d frames", st.name, st.bytes, st.frames)
}