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)
	}

	up := &stream{s: s, id: id, name: "c2g"}
	down := &stream{s: s, id: id, name: "g2c"}
	// 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(conn, upstream)
	}()
	wg.Wait()
	s.logf("#%d CONNECT %s closed after %s: %s, %s", id, target, since(start), up.summary(), down.summary())
}

// stream is one direction of a CONNECT, counting bytes and IMQ frames as they pass.
type stream struct {
	s      *Server
	id     int64
	name   string
	bytes  int64
	frames int
	other  bool // not IMQ framing
	reader imq.Reader
}

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.observe(buf[:n])
			if _, werr := dst.Write(buf[:n]); werr != nil {
				return
			}
		}
		if err != nil {
			return
		}
	}
}

func (st *stream) observe(data []byte) {
	st.bytes += int64(len(data))
	if st.other {
		return
	}
	frames, err := st.reader.Write(data)
	st.frames += len(frames)
	if st.s.Trace {
		for _, f := range frames {
			st.s.logf("#%d %s %s %dB", st.id, st.name, imq.TypeName(f.Type), f.Size)
		}
	}
	if err != nil {
		st.other = true
		st.s.logf("#%d %s isn't IMQ framing; passing it through", st.id, st.name)
	}
}

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)
}