diff --git a/DECISIONS.md b/DECISIONS.md index 833ec82..db54ab7 100644 --- a/DECISIONS.md +++ b/DECISIONS.md @@ -10,5 +10,5 @@ - 2026-02-25 m+git@andri.dk — Cache `CSSForDomain` results in `sync.Map` since the ruleset is immutable after loading. - 2026-02-25 m+git@andri.dk — Serve decompressed HTML to client after CSS injection (no re-compression). Proxy-to-client hop is typically localhost. - 2026-02-25 m+git@andri.dk — Go stdlib only. No third-party dependencies unless strongly warranted. -- 2026-02-25 m+git@andri.dk — WebSocket upgrade is not supported (Upgrade header stripped as hop-by-hop). Acceptable for an ad-blocking proxy. +- 2026-02-25 m+git@andri.dk — WebSocket upgrade supported for both ws:// and wss://. Upgrade headers are re-added after hop-by-hop stripping, then bidirectional copy bridges client and upstream after 101. - 2026-02-25 m+git@andri.dk — Cert cache has no eviction. Certs are generated with 24h validity. Acceptable for personal use. diff --git a/README.md b/README.md index 933f7df..aebe58e 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,6 @@ A proxy-server, that is capable of filtering ads from HTTPS/TLS traffic, using a ### Known Limitations -- **WebSocket**: `Upgrade` and `Connection` headers are stripped as hop-by-hop headers, preventing WebSocket connections through the proxy. - **Accept-Encoding downgrade**: For domains with element hiding (CSS injection) rules, `Accept-Encoding` is downgraded to `gzip` for all requests to that domain (not just HTML). Non-HTML resources lose brotli compression on those domains. - **No re-compression**: After decompressing gzip for CSS injection, HTML is served uncompressed to the client. This is fine when the proxy runs on localhost. - **Cert cache**: Generated TLS certificates are cached indefinitely with no eviction. Certificates have 24-hour validity but expired entries are never cleaned up. Fine for personal use. diff --git a/connect.go b/connect.go index 47c447b..79818bb 100644 --- a/connect.go +++ b/connect.go @@ -89,6 +89,8 @@ func (p *proxyHandler) proxyTLSRequests(clientTLS *tls.Conn, host, port string) continue } + upgradeReq := isWebSocketUpgrade(req.Header) + start := time.Now() req.URL.Scheme = "https" req.URL.Host = net.JoinHostPort(host, port) @@ -96,9 +98,15 @@ func (p *proxyHandler) proxyTLSRequests(clientTLS *tls.Conn, host, port string) removeHopByHopHeaders(req.Header) + // Re-add upgrade headers that were stripped as hop-by-hop + if upgradeReq { + req.Header.Set("Connection", "Upgrade") + req.Header.Set("Upgrade", "websocket") + } + // Downgrade to gzip-only when we may need to decompress HTML for CSS // injection. We can only decompress gzip, not brotli or other encodings. - if p.rules.CSSForDomain(host) != "" { + if !upgradeReq && p.rules.CSSForDomain(host) != "" { req.Header.Set("Accept-Encoding", "gzip") } @@ -108,6 +116,27 @@ func (p *proxyHandler) proxyTLSRequests(clientTLS *tls.Conn, host, port string) return } + // WebSocket upgrade: switch to bidirectional copy + if upgradeReq && resp.StatusCode == http.StatusSwitchingProtocols { + upstreamConn, ok := resp.Body.(io.ReadWriteCloser) + if !ok { + resp.Body.Close() + logError("connect/upgrade", io.ErrUnexpectedEOF) + return + } + defer upstreamConn.Close() + + resp.Body = nil + resp.Write(clientTLS) + + logRequest(req.Method, targetURL+" [websocket]", resp.StatusCode, time.Since(start)) + + // clientReader may have buffered bytes past the HTTP request, + // so we read from it (not raw clientTLS). Writes go to clientTLS. + bidirectionalCopy(&readerWriter{r: clientReader, w: clientTLS}, upstreamConn) + return + } + // Inject element hiding CSS into HTML responses (skip HEAD — no body to modify) if req.Method != http.MethodHead { if modified, ok := p.injectElementHidingCSS(resp, host); ok { diff --git a/http.go b/http.go index 5a7f38b..5e0d061 100644 --- a/http.go +++ b/http.go @@ -29,6 +29,12 @@ func (p *proxyHandler) handleHTTP(w http.ResponseWriter, r *http.Request) { return } + // WebSocket and other protocol upgrades need special handling + if isWebSocketUpgrade(r.Header) { + p.handleHTTPUpgrade(w, r) + return + } + start := time.Now() outReq, err := http.NewRequestWithContext(r.Context(), r.Method, r.URL.String(), r.Body) @@ -76,6 +82,74 @@ func (p *proxyHandler) handleHTTP(w http.ResponseWriter, r *http.Request) { logRequest(r.Method, r.URL.String(), resp.StatusCode, time.Since(start)) } +// handleHTTPUpgrade handles WebSocket and other protocol upgrade requests +// over plain HTTP. It preserves the upgrade headers, sends the request to +// the upstream, and if the upstream responds with 101, hijacks both sides +// and does bidirectional copy. +func (p *proxyHandler) handleHTTPUpgrade(w http.ResponseWriter, r *http.Request) { + outReq, err := http.NewRequestWithContext(r.Context(), r.Method, r.URL.String(), r.Body) + if err != nil { + logError("http/upgrade/new-request", err) + http.Error(w, "bad request", http.StatusBadRequest) + return + } + + copyHeaders(outReq.Header, r.Header) + removeHopByHopHeaders(outReq.Header) + + // Re-add the upgrade headers that were stripped as hop-by-hop + outReq.Header.Set("Connection", "Upgrade") + outReq.Header.Set("Upgrade", r.Header.Get("Upgrade")) + + resp, err := p.transport.RoundTrip(outReq) + if err != nil { + logError("http/upgrade/roundtrip", err) + http.Error(w, "upstream error", http.StatusBadGateway) + return + } + + if resp.StatusCode != http.StatusSwitchingProtocols { + // Not a 101 — fall back to normal response handling + defer resp.Body.Close() + copyHeaders(w.Header(), resp.Header) + removeHopByHopHeaders(w.Header()) + w.WriteHeader(resp.StatusCode) + io.Copy(w, resp.Body) + return + } + + // Get the raw upstream connection from the response body + upstreamConn, ok := resp.Body.(io.ReadWriteCloser) + if !ok { + resp.Body.Close() + http.Error(w, "upstream does not support hijacking", http.StatusInternalServerError) + return + } + defer upstreamConn.Close() + + // Hijack the client connection + hijacker, ok := w.(http.Hijacker) + if !ok { + http.Error(w, "hijacking not supported", http.StatusInternalServerError) + return + } + + clientConn, clientBuf, err := hijacker.Hijack() + if err != nil { + logError("http/upgrade/hijack", err) + return + } + defer clientConn.Close() + + // Write the 101 response to the client + resp.Body = nil // don't write the body, just headers + resp.Write(clientConn) + clientBuf.Flush() + + // Bidirectional copy between client and upstream + bidirectionalCopy(clientConn, upstreamConn) +} + // matchContextFromReferer extracts the page domain from the Referer header // for evaluating context-dependent filter options ($third-party, $domain). func matchContextFromReferer(referer string) blocklist.MatchContext { diff --git a/proxy_test.go b/proxy_test.go index 240487e..2bcb54f 100644 --- a/proxy_test.go +++ b/proxy_test.go @@ -1,10 +1,15 @@ package main import ( + "bufio" + "crypto/sha1" "crypto/tls" "crypto/x509" + "encoding/base64" "encoding/pem" + "fmt" "io" + "net" "net/http" "net/http/httptest" "net/url" @@ -75,6 +80,12 @@ func startTestEnv(t *testing.T, upstreamHandler http.Handler, rules *blocklist.R } } +// httpsHost returns the host:port of the HTTPS upstream server. +func (e *testEnv) httpsHost() string { + u, _ := url.Parse(e.httpsURL) + return u.Host +} + // httpClient returns an *http.Client configured to route through the proxy. // For HTTPS requests, it trusts the test CA. func (e *testEnv) httpClient(t *testing.T) *http.Client { @@ -812,3 +823,254 @@ func TestHeadRequestNoInjection(t *testing.T) { t.Errorf("status = %d, want %d", resp.StatusCode, http.StatusOK) } } + +// --- WebSocket test helpers --- + +// wsEchoHandler is a minimal WebSocket echo server using only stdlib. +// It performs the WebSocket handshake, then echoes back any frames it receives. +func wsEchoHandler(w http.ResponseWriter, r *http.Request) { + if !strings.EqualFold(r.Header.Get("Upgrade"), "websocket") { + http.Error(w, "not a websocket request", http.StatusBadRequest) + return + } + + // Compute Sec-WebSocket-Accept from the client key + key := r.Header.Get("Sec-WebSocket-Key") + acceptKey := computeWebSocketAccept(key) + + h := w.(http.Hijacker) + conn, buf, err := h.Hijack() + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + defer conn.Close() + + // Write the 101 response + buf.WriteString("HTTP/1.1 101 Switching Protocols\r\n") + buf.WriteString("Upgrade: websocket\r\n") + buf.WriteString("Connection: Upgrade\r\n") + buf.WriteString("Sec-WebSocket-Accept: " + acceptKey + "\r\n") + buf.WriteString("\r\n") + buf.Flush() + + // Echo loop: read a frame, write it back + for { + frame, err := readWSFrame(buf.Reader) + if err != nil { + return + } + writeWSFrame(conn, frame) + } +} + +// computeWebSocketAccept computes the Sec-WebSocket-Accept value per RFC 6455. +func computeWebSocketAccept(key string) string { + const websocketGUID = "258EAFA5-E914-47DA-95CA-5AB5DC76E45B" + h := sha1.New() + h.Write([]byte(key + websocketGUID)) + return base64.StdEncoding.EncodeToString(h.Sum(nil)) +} + +// wsFrame is a minimal WebSocket frame (text only, no masking on server side). +type wsFrame struct { + payload []byte +} + +// readWSFrame reads a single WebSocket frame. Handles client-masked frames. +func readWSFrame(r io.Reader) (wsFrame, error) { + header := make([]byte, 2) + if _, err := io.ReadFull(r, header); err != nil { + return wsFrame{}, err + } + + masked := header[1]&0x80 != 0 + length := int(header[1] & 0x7F) + + // Only support small frames for testing + if length == 126 || length == 127 { + return wsFrame{}, io.ErrUnexpectedEOF + } + + var maskKey [4]byte + if masked { + if _, err := io.ReadFull(r, maskKey[:]); err != nil { + return wsFrame{}, err + } + } + + payload := make([]byte, length) + if _, err := io.ReadFull(r, payload); err != nil { + return wsFrame{}, err + } + + if masked { + for i := range payload { + payload[i] ^= maskKey[i%4] + } + } + + return wsFrame{payload: payload}, nil +} + +// writeWSFrame writes a single unmasked text frame. +func writeWSFrame(w io.Writer, f wsFrame) error { + header := []byte{0x81, byte(len(f.payload))} // FIN + text opcode, length + if _, err := w.Write(header); err != nil { + return err + } + _, err := w.Write(f.payload) + return err +} + +// writeMaskedWSFrame writes a single masked text frame (required for client-to-server). +func writeMaskedWSFrame(w io.Writer, payload []byte) error { + maskKey := [4]byte{0x12, 0x34, 0x56, 0x78} // fixed mask for testing + masked := make([]byte, len(payload)) + for i := range payload { + masked[i] = payload[i] ^ maskKey[i%4] + } + header := []byte{0x81, byte(len(payload)) | 0x80} // FIN + text opcode, masked, length + if _, err := w.Write(header); err != nil { + return err + } + if _, err := w.Write(maskKey[:]); err != nil { + return err + } + _, err := w.Write(masked) + return err +} + +// dialWebSocketViaProxy performs a WebSocket handshake through the proxy. +// Returns the raw connection for sending/receiving frames. +func dialWebSocketViaProxy(proxyURL, targetURL string, tlsConfig *tls.Config) (net.Conn, error) { + parsed, _ := url.Parse(targetURL) + proxyParsed, _ := url.Parse(proxyURL) + + var conn net.Conn + var err error + + if parsed.Scheme == "wss" { + // For wss://, first CONNECT to establish the tunnel + conn, err = net.Dial("tcp", proxyParsed.Host) + if err != nil { + return nil, fmt.Errorf("dial proxy: %w", err) + } + + // Send CONNECT + host := parsed.Host + if !strings.Contains(host, ":") { + host += ":443" + } + fmt.Fprintf(conn, "CONNECT %s HTTP/1.1\r\nHost: %s\r\n\r\n", host, host) + + // Read CONNECT response + br := bufio.NewReader(conn) + resp, err := http.ReadResponse(br, nil) + if err != nil { + conn.Close() + return nil, fmt.Errorf("CONNECT response: %w", err) + } + if resp.StatusCode != http.StatusOK { + conn.Close() + return nil, fmt.Errorf("CONNECT status: %d", resp.StatusCode) + } + + // TLS handshake over the tunnel + tlsConn := tls.Client(conn, tlsConfig) + if err := tlsConn.Handshake(); err != nil { + conn.Close() + return nil, fmt.Errorf("TLS handshake: %w", err) + } + conn = tlsConn + } else { + // For ws://, connect to the proxy directly + conn, err = net.Dial("tcp", proxyParsed.Host) + if err != nil { + return nil, fmt.Errorf("dial proxy: %w", err) + } + } + + // Send WebSocket upgrade request. Over a CONNECT tunnel (wss://), use + // only the path since the TLS connection is already to the right host. + // Over plain HTTP proxy (ws://), use the full URL per HTTP proxy spec. + requestURI := targetURL + if parsed.Scheme == "wss" { + requestURI = parsed.RequestURI() + } + wsKey := base64.StdEncoding.EncodeToString([]byte("test-websocket-key!")) + reqStr := fmt.Sprintf( + "GET %s HTTP/1.1\r\nHost: %s\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Key: %s\r\nSec-WebSocket-Version: 13\r\n\r\n", + requestURI, parsed.Host, wsKey, + ) + if _, err := conn.Write([]byte(reqStr)); err != nil { + conn.Close() + return nil, fmt.Errorf("write upgrade: %w", err) + } + + // Read the 101 response + br := bufio.NewReader(conn) + resp, err := http.ReadResponse(br, nil) + if err != nil { + conn.Close() + return nil, fmt.Errorf("read upgrade response: %w", err) + } + if resp.StatusCode != http.StatusSwitchingProtocols { + conn.Close() + return nil, fmt.Errorf("expected 101, got %d", resp.StatusCode) + } + + return conn, nil +} + +func TestWebSocketHTTP(t *testing.T) { + env := startTestEnv(t, http.HandlerFunc(wsEchoHandler), nil) + + conn, err := dialWebSocketViaProxy(env.proxyURL, env.httpURL+"/ws", nil) + if err != nil { + t.Fatalf("WebSocket dial: %v", err) + } + defer conn.Close() + + // Send a message + msg := []byte("hello websocket") + if err := writeMaskedWSFrame(conn, msg); err != nil { + t.Fatalf("write frame: %v", err) + } + + // Read echo + frame, err := readWSFrame(conn) + if err != nil { + t.Fatalf("read frame: %v", err) + } + + if string(frame.payload) != "hello websocket" { + t.Errorf("echo = %q, want %q", frame.payload, "hello websocket") + } +} + +func TestWebSocketHTTPS(t *testing.T) { + env := startTestEnv(t, http.HandlerFunc(wsEchoHandler), nil) + + host, _, _ := net.SplitHostPort(env.httpsHost()) + tlsConfig := &tls.Config{RootCAs: env.caPool, ServerName: host} + conn, err := dialWebSocketViaProxy(env.proxyURL, "wss://"+env.httpsHost()+"/ws", tlsConfig) + if err != nil { + t.Fatalf("WebSocket dial: %v", err) + } + defer conn.Close() + + msg := []byte("hello secure websocket") + if err := writeMaskedWSFrame(conn, msg); err != nil { + t.Fatalf("write frame: %v", err) + } + + frame, err := readWSFrame(conn) + if err != nil { + t.Fatalf("read frame: %v", err) + } + + if string(frame.payload) != "hello secure websocket" { + t.Errorf("echo = %q, want %q", frame.payload, "hello secure websocket") + } +} diff --git a/websocket.go b/websocket.go new file mode 100644 index 0000000..d5a5efb --- /dev/null +++ b/websocket.go @@ -0,0 +1,39 @@ +package main + +import ( + "io" + "net/http" + "strings" +) + +// isWebSocketUpgrade returns true if the request is a WebSocket upgrade. +func isWebSocketUpgrade(h http.Header) bool { + return strings.EqualFold(h.Get("Upgrade"), "websocket") +} + +// readerWriter combines separate io.Reader and io.Writer into an +// io.ReadWriter. Used in CONNECT tunnels where reads come from a +// bufio.Reader (which may have buffered data) while writes go to +// the raw TLS connection. +type readerWriter struct { + r io.Reader + w io.Writer +} + +func (rw *readerWriter) Read(p []byte) (int, error) { return rw.r.Read(p) } +func (rw *readerWriter) Write(p []byte) (int, error) { return rw.w.Write(p) } + +// bidirectionalCopy copies data in both directions between a and b until +// one side closes or errors. Used for WebSocket and other protocol upgrades. +func bidirectionalCopy(a io.ReadWriter, b io.ReadWriter) { + done := make(chan struct{}, 1) + go func() { + io.Copy(a, b) + done <- struct{}{} + }() + go func() { + io.Copy(b, a) + done <- struct{}{} + }() + <-done +}