From 914fd83ac4d9e334f627b41f01124e1b5a9b8968 Mon Sep 17 00:00:00 2001 From: Jian Zeng Date: Sun, 20 Aug 2023 22:58:06 +0800 Subject: [PATCH] refactor: switch to shorter request id (#22) --- .dockerignore | 2 + Makefile | 4 +- cmd/client/{info.go => goreleaser.go} | 0 cmd/client/tcp.go | 9 ++-- cmd/client/udp.go | 9 ++-- cmd/client/utils.go | 1 + cmd/server/main.go | 14 +++--- cmd/server/tcp_test.go | 3 +- go.mod | 2 +- pkg/xnet/header.go | 67 ++++++++++++++++++++------- pkg/xnet/header_test.go | 21 +++++---- 11 files changed, 82 insertions(+), 50 deletions(-) create mode 100644 .dockerignore rename cmd/client/{info.go => goreleaser.go} (100%) diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..082be0b --- /dev/null +++ b/.dockerignore @@ -0,0 +1,2 @@ +dist/ +.git/ diff --git a/Makefile b/Makefile index d8361e0..4bbd471 100644 --- a/Makefile +++ b/Makefile @@ -1,9 +1,7 @@ IMAGE_TAG ?= latest NAME := kubectl-relay -VERSION ?= $(shell git describe --tags || echo "unknown") -GO_LDFLAGS = "-w -s -X github.com/knight42/krelay/pkg/constants.ClientVersion=$(VERSION)" -GOBUILD = CGO_ENABLED=0 go build -trimpath -ldflags $(GO_LDFLAGS) +GOBUILD = CGO_ENABLED=0 go build -trimpath .PHONY: server-image push-server-image server-image: diff --git a/cmd/client/info.go b/cmd/client/goreleaser.go similarity index 100% rename from cmd/client/info.go rename to cmd/client/goreleaser.go diff --git a/cmd/client/tcp.go b/cmd/client/tcp.go index 1c2a164..b280979 100644 --- a/cmd/client/tcp.go +++ b/cmd/client/tcp.go @@ -4,7 +4,6 @@ import ( "io" "net" - "github.com/google/uuid" "k8s.io/apimachinery/pkg/util/httpstream" "k8s.io/klog/v2" @@ -16,17 +15,17 @@ import ( func handleTCPConn(clientConn net.Conn, serverConn httpstream.Connection, dstAddr xnet.Addr, dstPort uint16) { defer clientConn.Close() - requestID := uuid.New() - kvs := []any{constants.LogFieldRequestID, requestID.String()} + requestID := xnet.NewRequestID() + kvs := []any{constants.LogFieldRequestID, requestID} defer klog.V(4).InfoS("handleTCPConn exit", kvs...) klog.InfoS("Handling tcp connection", - constants.LogFieldRequestID, requestID.String(), + constants.LogFieldRequestID, requestID, constants.LogFieldDestAddr, xnet.JoinHostPort(dstAddr.String(), dstPort), constants.LogFieldLocalAddr, clientConn.LocalAddr().String(), "clientAddr", clientConn.RemoteAddr().String(), ) - dataStream, errorChan, err := createStream(serverConn, requestID.String()) + dataStream, errorChan, err := createStream(serverConn, requestID) if err != nil { klog.ErrorS(err, "Fail to create stream", kvs...) return diff --git a/cmd/client/udp.go b/cmd/client/udp.go index 0ad7976..1f2a1cb 100644 --- a/cmd/client/udp.go +++ b/cmd/client/udp.go @@ -3,7 +3,6 @@ package main import ( "net" - "github.com/google/uuid" "k8s.io/apimachinery/pkg/util/httpstream" "k8s.io/klog/v2" @@ -13,20 +12,20 @@ import ( ) func handleUDPConn(clientConn net.PacketConn, cliAddr net.Addr, dataCh chan []byte, finish chan<- string, serverConn httpstream.Connection, dstAddr xnet.Addr, dstPort uint16) { - requestID := uuid.New() - kvs := []any{constants.LogFieldDestAddr, requestID.String()} + requestID := xnet.NewRequestID() + kvs := []any{constants.LogFieldDestAddr, requestID} defer klog.V(4).InfoS("handleUDPConn exit", kvs...) defer func() { finish <- cliAddr.String() }() klog.InfoS("Handling udp connection", - constants.LogFieldRequestID, requestID.String(), + constants.LogFieldRequestID, requestID, constants.LogFieldDestAddr, xnet.JoinHostPort(dstAddr.String(), dstPort), constants.LogFieldLocalAddr, clientConn.LocalAddr().String(), "clientAddr", cliAddr.String(), ) - dataStream, errorChan, err := createStream(serverConn, requestID.String()) + dataStream, errorChan, err := createStream(serverConn, requestID) if err != nil { klog.ErrorS(err, "Fail to create stream", kvs...) return diff --git a/cmd/client/utils.go b/cmd/client/utils.go index 4685b9f..7b87ef4 100644 --- a/cmd/client/utils.go +++ b/cmd/client/utils.go @@ -36,6 +36,7 @@ func ensureServerPod(ctx context.Context, cs kubernetes.Interface, svrImg, names GenerateName: constants.ServerName + "-", Labels: map[string]string{ "app.kubernetes.io/name": constants.ServerName, + "app": constants.ServerName, }, }, Spec: corev1.PodSpec{ diff --git a/cmd/server/main.go b/cmd/server/main.go index 46bcadf..0887768 100644 --- a/cmd/server/main.go +++ b/cmd/server/main.go @@ -56,24 +56,24 @@ func handleConn(ctx context.Context, c *net.TCPConn, dialer *net.Dialer) { case xnet.ProtocolTCP: upstreamConn, err := dialer.DialContext(ctx, constants.ProtocolTCP, dstAddr) if err != nil { - klog.ErrorS(err, "Fail to create tcp connection", constants.LogFieldRequestID, hdr.RequestID.String(), constants.LogFieldDestAddr, dstAddr) + klog.ErrorS(err, "Fail to create tcp connection", constants.LogFieldRequestID, hdr.RequestID, constants.LogFieldDestAddr, dstAddr) return } - klog.InfoS("Start proxy tcp request", constants.LogFieldRequestID, hdr.RequestID.String(), constants.LogFieldDestAddr, dstAddr) - xnet.ProxyTCP(hdr.RequestID.String(), c, upstreamConn.(*net.TCPConn)) + klog.InfoS("Start proxy tcp request", constants.LogFieldRequestID, hdr.RequestID, constants.LogFieldDestAddr, dstAddr) + xnet.ProxyTCP(hdr.RequestID, c, upstreamConn.(*net.TCPConn)) case xnet.ProtocolUDP: upstreamConn, err := dialer.DialContext(ctx, constants.ProtocolUDP, dstAddr) if err != nil { - klog.ErrorS(err, "Fail to create udp connection", constants.LogFieldRequestID, hdr.RequestID.String(), constants.LogFieldDestAddr, dstAddr) + klog.ErrorS(err, "Fail to create udp connection", constants.LogFieldRequestID, hdr.RequestID, constants.LogFieldDestAddr, dstAddr) return } - klog.InfoS("Start proxy udp request", constants.LogFieldRequestID, hdr.RequestID.String(), constants.LogFieldDestAddr, dstAddr) + klog.InfoS("Start proxy udp request", constants.LogFieldRequestID, hdr.RequestID, constants.LogFieldDestAddr, dstAddr) udpConn := &xnet.UDPConn{UDPConn: upstreamConn.(*net.UDPConn)} - xnet.ProxyUDP(hdr.RequestID.String(), c, udpConn) + xnet.ProxyUDP(hdr.RequestID, c, udpConn) default: - klog.InfoS("Unknown protocol", constants.LogFieldRequestID, hdr.RequestID.String(), constants.LogFieldDestAddr, dstAddr, constants.LogFieldProtocol, hdr.Protocol) + klog.InfoS("Unknown protocol", constants.LogFieldRequestID, hdr.RequestID, constants.LogFieldDestAddr, dstAddr, constants.LogFieldProtocol, hdr.Protocol) } } diff --git a/cmd/server/tcp_test.go b/cmd/server/tcp_test.go index 1be070e..744c6b2 100644 --- a/cmd/server/tcp_test.go +++ b/cmd/server/tcp_test.go @@ -9,7 +9,6 @@ import ( "testing" "time" - "github.com/google/uuid" "github.com/stretchr/testify/require" "github.com/knight42/krelay/pkg/testutils/tcp" @@ -41,7 +40,7 @@ func TestHandleTCPConn(t *testing.T) { return nil, fmt.Errorf("dial: %w", err) } hdr := xnet.Header{ - RequestID: uuid.New(), + RequestID: xnet.NewRequestID(), Protocol: xnet.ProtocolTCP, Port: tsPort, Addr: xnet.AddrFromHost(tsURL.Hostname()), diff --git a/go.mod b/go.mod index 6b59203..7b8f62b 100644 --- a/go.mod +++ b/go.mod @@ -3,7 +3,6 @@ module github.com/knight42/krelay go 1.21 require ( - github.com/google/uuid v1.3.0 github.com/spf13/cobra v1.7.0 github.com/stretchr/testify v1.8.4 k8s.io/api v0.28.0 @@ -29,6 +28,7 @@ require ( github.com/google/go-cmp v0.5.9 // indirect github.com/google/gofuzz v1.2.0 // indirect github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510 // indirect + github.com/google/uuid v1.3.0 // indirect github.com/gregjones/httpcache v0.0.0-20190611155906-901d90724c79 // indirect github.com/imdario/mergo v0.3.16 // indirect github.com/inconshreveable/mousetrap v1.1.0 // indirect diff --git a/pkg/xnet/header.go b/pkg/xnet/header.go index b799c50..7747ebe 100644 --- a/pkg/xnet/header.go +++ b/pkg/xnet/header.go @@ -4,15 +4,27 @@ import ( "encoding/binary" "fmt" "io" + "math/rand" +) - "github.com/google/uuid" +const ( + lengthAllMandatoryFields = 12 // 1(version) + 2(total length) + 5(request id) + 1(protocol) + 2(port) + 1(addr type) + lengthRequestID = 5 ) -const totalLenOfOtherFields = 23 // 1(version) + 2(total length) + 16(uuid) + 1(protocol) + 2(port) + 1(addr type) +var letterRunes = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789") + +func NewRequestID() string { + b := make([]rune, lengthRequestID) + for i := range b { + b[i] = letterRunes[rand.Intn(len(letterRunes))] + } + return string(b) +} type Header struct { Version byte - RequestID uuid.UUID + RequestID string Protocol byte Port uint16 Addr Addr @@ -20,15 +32,29 @@ type Header struct { func (h *Header) Marshal() []byte { addrBytes := h.Addr.Marshal() - totalLen := totalLenOfOtherFields + len(addrBytes) + totalLen := lengthAllMandatoryFields + len(addrBytes) buf := make([]byte, totalLen) - buf[0] = h.Version - binary.BigEndian.PutUint16(buf[1:3], uint16(totalLen)) - copy(buf[3:19], h.RequestID[:]) - buf[19] = h.Protocol - binary.BigEndian.PutUint16(buf[20:22], h.Port) - buf[22] = h.Addr.typ - copy(buf[23:], addrBytes) + + cursor := 0 + buf[cursor] = h.Version + cursor++ + + binary.BigEndian.PutUint16(buf[cursor:cursor+2], uint16(totalLen)) + cursor += 2 + + copy(buf[cursor:cursor+lengthRequestID], h.RequestID) + cursor += lengthRequestID + + buf[cursor] = h.Protocol + cursor++ + + binary.BigEndian.PutUint16(buf[cursor:cursor+2], h.Port) + cursor += 2 + + buf[cursor] = h.Addr.typ + cursor++ + + copy(buf[cursor:], addrBytes) return buf } @@ -41,7 +67,7 @@ func (h *Header) FromReader(r io.Reader) error { h.Version = lengthBuf[0] // TODO: handle different versions totalLen := binary.BigEndian.Uint16(lengthBuf[1:]) - if totalLen < totalLenOfOtherFields { + if totalLen < lengthAllMandatoryFields { return fmt.Errorf("body too short: %d", totalLen) } @@ -51,12 +77,19 @@ func (h *Header) FromReader(r io.Reader) error { return fmt.Errorf("read full body: %w", err) } - reqIDBytes := bodyBuf[:16] - proto := bodyBuf[16] - port := binary.BigEndian.Uint16(bodyBuf[17:19]) - h.RequestID, _ = uuid.FromBytes(reqIDBytes) + cursor := 0 + reqIDBytes := bodyBuf[cursor : cursor+lengthRequestID] + cursor += lengthRequestID + + proto := bodyBuf[cursor] + cursor++ + + port := binary.BigEndian.Uint16(bodyBuf[cursor : cursor+2]) + cursor += 2 + + h.RequestID = string(reqIDBytes) h.Protocol = proto h.Port = port - h.Addr = AddrFromBytes(bodyBuf[19], bodyBuf[20:]) + h.Addr = AddrFromBytes(bodyBuf[cursor], bodyBuf[cursor+1:]) return nil } diff --git a/pkg/xnet/header_test.go b/pkg/xnet/header_test.go index 0339cfe..fb80bc4 100644 --- a/pkg/xnet/header_test.go +++ b/pkg/xnet/header_test.go @@ -5,10 +5,11 @@ import ( "net" "testing" - "github.com/google/uuid" "github.com/stretchr/testify/require" ) +var fakeRequestID = "00000" + var headerCases = map[string]struct { hdr Header bytes []byte @@ -16,15 +17,15 @@ var headerCases = map[string]struct { "host": { hdr: Header{ Version: 1, - RequestID: uuid.UUID{}, + RequestID: fakeRequestID, Protocol: ProtocolTCP, Port: 80, Addr: AddrFromHost("a.com"), }, bytes: []byte{ 1, - 0, 0x1c, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0x11, + 0x30, 0x30, 0x30, 0x30, 0x30, 0, 0, 80, 1, @@ -34,15 +35,15 @@ var headerCases = map[string]struct { "ipv4": { hdr: Header{ Version: 0, - RequestID: uuid.UUID{}, + RequestID: fakeRequestID, Protocol: ProtocolUDP, Port: 53, Addr: AddrFromBytes(AddrTypeIP, net.IPv4(192, 168, 1, 1).To4()), }, bytes: []byte{ 0, - 0, 0x1b, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0x10, + 0x30, 0x30, 0x30, 0x30, 0x30, 1, 0, 53, 0, @@ -52,15 +53,15 @@ var headerCases = map[string]struct { "ipv6": { hdr: Header{ Version: 0, - RequestID: uuid.UUID{}, + RequestID: fakeRequestID, Protocol: ProtocolTCP, Port: 8080, Addr: AddrFromBytes(AddrTypeIP, net.IPv6loopback), }, bytes: []byte{ 0, - 0, 0x27, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0x1c, + 0x30, 0x30, 0x30, 0x30, 0x30, 0, 0x1f, 0x90, 0,