tcp/udp fix

Dockerfile

@@ -1,20 +1,30 @@
-FROM golang:1.22.1 as builder
+FROM alpine:latest AS codemirror
+RUN apk add --no-cache unzip wget make
+COPY Makefile .
+RUN make setup-codemirror
 
-COPY go.mod /app/go.mod
-COPY src/ /app/src
-COPY Makefile /app
-WORKDIR /app
-RUN make get
-RUN make build
+FROM golang:1.22.1-alpine as builder
+COPY src/ /src
+COPY go.mod go.sum /src/go-proxy
+WORKDIR /src/go-proxy
+RUN --mount=type=cache,target="/go/pkg/mod" \
+    go mod download
+
+ENV GOCACHE=/root/.cache/go-build
+RUN --mount=type=cache,target="/go/pkg/mod" \
+    --mount=type=cache,target="/root/.cache/go-build" \
+    CGO_ENABLED=0 GOOS=linux go build -pgo=auto -o go-proxy
 
 FROM alpine:latest
 
 LABEL maintainer="[email protected]"
 
 RUN apk add --no-cache tzdata
-COPY --from=builder /app/bin/go-proxy /app/
+RUN mkdir -p /app/templates
+COPY --from=codemirror templates/codemirror/ /app/templates/codemirror
 COPY templates/ /app/templates
 COPY schema/ /app/schema
+COPY --from=builder /src/go-proxy /app/
 
 RUN chmod +x /app/go-proxy
 ENV DOCKER_HOST unix:///var/run/docker.sock

Makefile

@@ -11,6 +11,7 @@ setup-codemirror:
 	wget https://codemirror.net/5/codemirror.zip
 	unzip codemirror.zip
 	rm codemirror.zip
+	mkdir -p templates
 	mv codemirror-* templates/codemirror
 
 build:
@@ -35,6 +36,6 @@ udp-server:
 		-p 9999:9999/udp \
 		--label proxy.test-udp.scheme=udp \
 		--label proxy.test-udp.port=20003:9999 \
-		--network data_default \
+		--network host \
 		--name test-udp \
 		$$(docker build -q -f udp-test-server.Dockerfile .)

src/go-proxy/io.go

@@ -2,13 +2,37 @@ package main
 
 import (
 	"context"
+	"errors"
+	"fmt"
 	"io"
-	"sync"
+	"os"
+	"sync/atomic"
 )
 
+type Reader interface {
+	Read() ([]byte, error)
+}
+
+type FileReader struct {
+	Path string
+}
+
+func (r *FileReader) Read() ([]byte, error) {
+	return os.ReadFile(r.Path)
+}
+
+type ByteReader struct {
+	Data []byte
+}
+
+func (r *ByteReader) Read() ([]byte, error) {
+	return r.Data, nil
+}
+
 type ReadCloser struct {
-	ctx context.Context
-	r   io.ReadCloser
+	ctx    context.Context
+	r      io.ReadCloser
+	closed atomic.Bool
 }
 
 func (r *ReadCloser) Read(p []byte) (int, error) {
@@ -21,46 +45,41 @@ func (r *ReadCloser) Read(p []byte) (int, error) {
 }
 
 func (r *ReadCloser) Close() error {
+	if r.closed.Load() {
+		return nil
+	}
+	r.closed.Store(true)
 	return r.r.Close()
 }
 
 type Pipe struct {
 	r      ReadCloser
 	w      io.WriteCloser
-	wg     sync.WaitGroup
 	ctx    context.Context
 	cancel context.CancelFunc
 }
 
 func NewPipe(ctx context.Context, r io.ReadCloser, w io.WriteCloser) *Pipe {
 	ctx, cancel := context.WithCancel(ctx)
 	return &Pipe{
-		r:      ReadCloser{ctx, r},
+		r:      ReadCloser{ctx: ctx, r: r},
 		w:      w,
 		ctx:    ctx,
 		cancel: cancel,
 	}
 }
 
-func (p *Pipe) Start() {
-	p.wg.Add(1)
-	go func() {
-		Copy(p.ctx, p.w, &p.r)
-		p.wg.Done()
-	}()
+func (p *Pipe) Start() error {
+	return Copy(p.ctx, p.w, &p.r)
 }
 
-func (p *Pipe) Stop() {
+func (p *Pipe) Stop() error {
 	p.cancel()
-	p.wg.Wait()
-}
-
-func (p *Pipe) Close() (error, error) {
-	return p.r.Close(), p.w.Close()
+	return errors.Join(fmt.Errorf("read: %w", p.r.Close()), fmt.Errorf("write: %w", p.w.Close()))
 }
 
-func (p *Pipe) Wait() {
-	p.wg.Wait()
+func (p *Pipe) Write(b []byte) (int, error) {
+	return p.w.Write(b)
 }
 
 type BidirectionalPipe struct {
@@ -75,26 +94,34 @@ func NewBidirectionalPipe(ctx context.Context, rw1 io.ReadWriteCloser, rw2 io.Re
 	}
 }
 
-func (p *BidirectionalPipe) Start() {
-	p.pSrcDst.Start()
-	p.pDstSrc.Start()
-}
-
-func (p *BidirectionalPipe) Stop() {
-	p.pSrcDst.Stop()
-	p.pDstSrc.Stop()
+func NewBidirectionalPipeIntermediate(ctx context.Context, listener io.ReadCloser, client io.ReadWriteCloser, target io.ReadWriteCloser) *BidirectionalPipe {
+	return &BidirectionalPipe{
+		pSrcDst: *NewPipe(ctx, listener, client),
+		pDstSrc: *NewPipe(ctx, client, target),
+	}
 }
 
-func (p *BidirectionalPipe) Close() (error, error) {
-	return p.pSrcDst.Close()
+func (p *BidirectionalPipe) Start() error {
+	errCh := make(chan error, 2)
+	go func() {
+		errCh <- p.pSrcDst.Start()
+	}()
+	go func() {
+		errCh <- p.pDstSrc.Start()
+	}()
+	for err := range errCh {
+		if err != nil {
+			return err
+		}
+	}
+	return nil
 }
 
-func (p *BidirectionalPipe) Wait() {
-	p.pSrcDst.Wait()
-	p.pDstSrc.Wait()
+func (p *BidirectionalPipe) Stop() error {
+	return errors.Join(p.pSrcDst.Stop(), p.pDstSrc.Stop())
 }
 
 func Copy(ctx context.Context, dst io.WriteCloser, src io.ReadCloser) error {
-	_, err := io.Copy(dst, &ReadCloser{ctx, src})
+	_, err := io.Copy(dst, &ReadCloser{ctx: ctx, r: src})
 	return err
 }

src/go-proxy/route.go

@@ -15,7 +15,6 @@ func NewRoute(cfg *ProxyConfig) (Route, error) {
 		if err != nil {
 			return nil, NewNestedErrorFrom(err).Subject(cfg.Alias)
 		}
-		streamRoutes.Set(id, route)
 		return route, nil
 	} else {
 		httpRoutes.Ensure(cfg.Alias)

src/go-proxy/stream_route.go

@@ -143,6 +143,7 @@ func (route *StreamRouteBase) Start() {
 		route.l.Errorf("failed to setup: %v", err)
 		return
 	}
+	streamRoutes.Set(route.id, route)
 	route.started = true
 	route.wg.Add(2)
 	go route.grAcceptConnections()

src/go-proxy/tcp_route.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"net"
+	"sync"
 	"time"
 )
 
@@ -14,12 +15,15 @@ type Pipes []*BidirectionalPipe
 type TCPRoute struct {
 	*StreamRouteBase
 	listener net.Listener
+	pipe     Pipes
+	mu       sync.Mutex
 }
 
 func NewTCPRoute(base *StreamRouteBase) StreamImpl {
 	return &TCPRoute{
 		StreamRouteBase: base,
 		listener:        nil,
+		pipe:            make(Pipes, 0),
 	}
 }
 
@@ -40,7 +44,6 @@ func (route *TCPRoute) Handle(c interface{}) error {
 	clientConn := c.(net.Conn)
 
 	defer clientConn.Close()
-	defer route.wg.Done()
 
 	ctx, cancel := context.WithTimeout(context.Background(), tcpDialTimeout)
 	defer cancel()
@@ -58,11 +61,12 @@ func (route *TCPRoute) Handle(c interface{}) error {
 		<-route.stopCh
 		pipeCancel()
 	}()
+
+	route.mu.Lock()
 	pipe := NewBidirectionalPipe(pipeCtx, clientConn, serverConn)
-	pipe.Start()
-	pipe.Wait()
-	pipe.Close()
-	return nil
+	route.pipe = append(route.pipe, pipe)
+	route.mu.Unlock()
+	return pipe.Start()
 }
 
 func (route *TCPRoute) CloseListeners() {
@@ -71,4 +75,9 @@ func (route *TCPRoute) CloseListeners() {
 	}
 	route.listener.Close()
 	route.listener = nil
+	for _, pipe := range route.pipe {
+		if err := pipe.Stop(); err != nil {
+			route.l.Error(err)
+		}
+	}
 }