1package tinygit23import (4 "bufio"5 "context"6 "crypto/rand"7 "crypto/subtle"8 "encoding/base64"9 "encoding/hex"10 "errors"11 "io"12 "net"13 "net/http"14 "strings"15 "sync"16 "sync/atomic"17)1819var errRepairByteLimit = errors.New("PR repair transfer exceeds byte limit")2021// A single budget covers every connection, ref and source in one repair.22// Counting tunneled TLS bytes also bounds HTTPS without terminating its TLS.23type byteBudget struct{ used atomic.Int64 }2425func (b *byteBudget) take(wanted, limit int64) int64 {26 for {27 old := b.used.Load()28 allowed := min(wanted, max(0, limit-old))29 if b.used.CompareAndSwap(old, old+allowed) {30 return allowed31 }32 }33}3435type repairProxy struct {36 listener net.Listener37 // credential is the Proxy-Authorization value Git presents. The listener38 // is loopback only, but other local processes cannot borrow the tunnel.39 credential string40 url string41 ctx context.Context42 cancel context.CancelFunc43 budget *byteBudget44 limit int6445 targetIP string46 targetHost string47 targetPort string48 mu sync.Mutex49 connections map[net.Conn]struct{}50 closeOnce sync.Once51}5253func newRepairProxy(ctx context.Context, targetIP, targetHost, targetPort string, budget *byteBudget, limit int64) (*repairProxy, error) {54 listener, err := net.Listen("tcp", "127.0.0.1:0")55 if err != nil {56 return nil, err57 }58 var secret [16]byte59 if _, err := rand.Read(secret[:]); err != nil {60 _ = listener.Close()61 return nil, err62 }63 proxyCtx, cancel := context.WithCancel(ctx)64 token := hex.EncodeToString(secret[:])65 credential := "Basic " + base64.StdEncoding.EncodeToString([]byte("tiny:"+token))66 p := &repairProxy{listener: listener, credential: credential, ctx: proxyCtx, cancel: cancel, budget: budget, limit: limit, targetIP: targetIP, targetHost: targetHost, targetPort: targetPort, connections: make(map[net.Conn]struct{})}67 p.url = "http://tiny:" + token + "@" + listener.Addr().String()68 context.AfterFunc(proxyCtx, func() { _ = p.Close() })69 go p.serve()70 return p, nil71}7273func (p *repairProxy) URL() string { return p.url }7475func (p *repairProxy) Close() error {76 p.closeOnce.Do(func() {77 p.cancel()78 _ = p.listener.Close()79 p.mu.Lock()80 defer p.mu.Unlock()81 for conn := range p.connections {82 _ = conn.Close()83 }84 })85 return nil86}8788func (p *repairProxy) track(conn net.Conn) bool {89 p.mu.Lock()90 defer p.mu.Unlock()91 if p.ctx.Err() != nil {92 _ = conn.Close()93 return false94 }95 p.connections[conn] = struct{}{}96 return true97}9899func (p *repairProxy) release(conn net.Conn) {100 _ = conn.Close()101 p.mu.Lock()102 delete(p.connections, conn)103 p.mu.Unlock()104}105106func (p *repairProxy) serve() {107 for {108 conn, err := p.listener.Accept()109 if err != nil {110 return111 }112 if p.track(conn) {113 go p.handle(conn)114 }115 }116}117118func (p *repairProxy) handle(client net.Conn) {119 defer p.release(client)120 reader := bufio.NewReader(client)121 request, err := http.ReadRequest(reader)122 if err != nil {123 return124 }125 if request.Method != http.MethodConnect && request.Method != http.MethodGet && request.Method != http.MethodPost {126 return127 }128 if subtle.ConstantTimeCompare([]byte(request.Header.Get("Proxy-Authorization")), []byte(p.credential)) != 1 {129 _, _ = io.WriteString(client, "HTTP/1.1 407 Proxy Authentication Required\r\nProxy-Authenticate: Basic realm=\"tiny\"\r\nContent-Length: 0\r\nConnection: close\r\n\r\n")130 return131 }132 host := request.URL.Hostname()133 port := request.URL.Port()134 if port == "" {135 port = "80"136 if request.Method == http.MethodConnect || request.URL.Scheme == "https" {137 port = "443"138 }139 }140 if !strings.EqualFold(host, p.targetHost) || port != p.targetPort {141 return142 }143 if p.budget.used.Load() >= p.limit {144 return145 }146 // Dial only the already validated address. The proxy never resolves a147 // host supplied by the client and never follows HTTP redirects.148 upstream, err := (&net.Dialer{}).DialContext(p.ctx, "tcp", net.JoinHostPort(p.targetIP, p.targetPort))149 if err != nil {150 return151 }152 if !p.track(upstream) {153 return154 }155 defer p.release(upstream)156 if request.Method == http.MethodConnect {157 if _, err := io.WriteString(repairBudgetWriter{p, client}, "HTTP/1.1 200 Connection Established\r\n\r\n"); err != nil {158 return159 }160 done := make(chan struct{}, 2)161 go func() { _, _ = io.Copy(repairBudgetWriter{p, client}, upstream); done <- struct{}{} }()162 // Keep bytes already buffered while reading CONNECT (e.g., a TLS163 // handshake sent in the same write as the request).164 go func() { _, _ = io.Copy(repairBudgetWriter{p, upstream}, reader); done <- struct{}{} }()165 select {166 case <-done:167 case <-p.ctx.Done():168 }169 return170 }171 request.RequestURI = ""172 request.Close = true173 if err := request.Write(repairBudgetWriter{p, upstream}); err != nil {174 return175 }176 _, _ = io.Copy(repairBudgetWriter{p, client}, upstream)177}178179type repairBudgetWriter struct {180 proxy *repairProxy181 dst io.Writer182}183184func (w repairBudgetWriter) Write(data []byte) (int, error) {185 allowed := w.proxy.budget.take(int64(len(data)), w.proxy.limit)186 n, err := w.dst.Write(data[:allowed])187 if err == nil && allowed < int64(len(data)) {188 err = errRepairByteLimit189 }190 return n, err191}