Skip to content

Commit 914fd83

Browse files
authored
refactor: switch to shorter request id (#22)
1 parent 3d8f08c commit 914fd83

File tree

11 files changed

+82
-50
lines changed

11 files changed

+82
-50
lines changed

.dockerignore

+2
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
dist/
2+
.git/

Makefile

+1-3
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,7 @@
11
IMAGE_TAG ?= latest
22

33
NAME := kubectl-relay
4-
VERSION ?= $(shell git describe --tags || echo "unknown")
5-
GO_LDFLAGS = "-w -s -X github.com/knight42/krelay/pkg/constants.ClientVersion=$(VERSION)"
6-
GOBUILD = CGO_ENABLED=0 go build -trimpath -ldflags $(GO_LDFLAGS)
4+
GOBUILD = CGO_ENABLED=0 go build -trimpath
75

86
.PHONY: server-image push-server-image
97
server-image:
File renamed without changes.

cmd/client/tcp.go

+4-5
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@ import (
44
"io"
55
"net"
66

7-
"github.com/google/uuid"
87
"k8s.io/apimachinery/pkg/util/httpstream"
98
"k8s.io/klog/v2"
109

@@ -16,17 +15,17 @@ import (
1615
func handleTCPConn(clientConn net.Conn, serverConn httpstream.Connection, dstAddr xnet.Addr, dstPort uint16) {
1716
defer clientConn.Close()
1817

19-
requestID := uuid.New()
20-
kvs := []any{constants.LogFieldRequestID, requestID.String()}
18+
requestID := xnet.NewRequestID()
19+
kvs := []any{constants.LogFieldRequestID, requestID}
2120
defer klog.V(4).InfoS("handleTCPConn exit", kvs...)
2221
klog.InfoS("Handling tcp connection",
23-
constants.LogFieldRequestID, requestID.String(),
22+
constants.LogFieldRequestID, requestID,
2423
constants.LogFieldDestAddr, xnet.JoinHostPort(dstAddr.String(), dstPort),
2524
constants.LogFieldLocalAddr, clientConn.LocalAddr().String(),
2625
"clientAddr", clientConn.RemoteAddr().String(),
2726
)
2827

29-
dataStream, errorChan, err := createStream(serverConn, requestID.String())
28+
dataStream, errorChan, err := createStream(serverConn, requestID)
3029
if err != nil {
3130
klog.ErrorS(err, "Fail to create stream", kvs...)
3231
return

cmd/client/udp.go

+4-5
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@ package main
33
import (
44
"net"
55

6-
"github.com/google/uuid"
76
"k8s.io/apimachinery/pkg/util/httpstream"
87
"k8s.io/klog/v2"
98

@@ -13,20 +12,20 @@ import (
1312
)
1413

1514
func handleUDPConn(clientConn net.PacketConn, cliAddr net.Addr, dataCh chan []byte, finish chan<- string, serverConn httpstream.Connection, dstAddr xnet.Addr, dstPort uint16) {
16-
requestID := uuid.New()
17-
kvs := []any{constants.LogFieldDestAddr, requestID.String()}
15+
requestID := xnet.NewRequestID()
16+
kvs := []any{constants.LogFieldDestAddr, requestID}
1817
defer klog.V(4).InfoS("handleUDPConn exit", kvs...)
1918
defer func() {
2019
finish <- cliAddr.String()
2120
}()
2221
klog.InfoS("Handling udp connection",
23-
constants.LogFieldRequestID, requestID.String(),
22+
constants.LogFieldRequestID, requestID,
2423
constants.LogFieldDestAddr, xnet.JoinHostPort(dstAddr.String(), dstPort),
2524
constants.LogFieldLocalAddr, clientConn.LocalAddr().String(),
2625
"clientAddr", cliAddr.String(),
2726
)
2827

29-
dataStream, errorChan, err := createStream(serverConn, requestID.String())
28+
dataStream, errorChan, err := createStream(serverConn, requestID)
3029
if err != nil {
3130
klog.ErrorS(err, "Fail to create stream", kvs...)
3231
return

cmd/client/utils.go

+1
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ func ensureServerPod(ctx context.Context, cs kubernetes.Interface, svrImg, names
3636
GenerateName: constants.ServerName + "-",
3737
Labels: map[string]string{
3838
"app.kubernetes.io/name": constants.ServerName,
39+
"app": constants.ServerName,
3940
},
4041
},
4142
Spec: corev1.PodSpec{

cmd/server/main.go

+7-7
Original file line numberDiff line numberDiff line change
@@ -56,24 +56,24 @@ func handleConn(ctx context.Context, c *net.TCPConn, dialer *net.Dialer) {
5656
case xnet.ProtocolTCP:
5757
upstreamConn, err := dialer.DialContext(ctx, constants.ProtocolTCP, dstAddr)
5858
if err != nil {
59-
klog.ErrorS(err, "Fail to create tcp connection", constants.LogFieldRequestID, hdr.RequestID.String(), constants.LogFieldDestAddr, dstAddr)
59+
klog.ErrorS(err, "Fail to create tcp connection", constants.LogFieldRequestID, hdr.RequestID, constants.LogFieldDestAddr, dstAddr)
6060
return
6161
}
62-
klog.InfoS("Start proxy tcp request", constants.LogFieldRequestID, hdr.RequestID.String(), constants.LogFieldDestAddr, dstAddr)
63-
xnet.ProxyTCP(hdr.RequestID.String(), c, upstreamConn.(*net.TCPConn))
62+
klog.InfoS("Start proxy tcp request", constants.LogFieldRequestID, hdr.RequestID, constants.LogFieldDestAddr, dstAddr)
63+
xnet.ProxyTCP(hdr.RequestID, c, upstreamConn.(*net.TCPConn))
6464

6565
case xnet.ProtocolUDP:
6666
upstreamConn, err := dialer.DialContext(ctx, constants.ProtocolUDP, dstAddr)
6767
if err != nil {
68-
klog.ErrorS(err, "Fail to create udp connection", constants.LogFieldRequestID, hdr.RequestID.String(), constants.LogFieldDestAddr, dstAddr)
68+
klog.ErrorS(err, "Fail to create udp connection", constants.LogFieldRequestID, hdr.RequestID, constants.LogFieldDestAddr, dstAddr)
6969
return
7070
}
71-
klog.InfoS("Start proxy udp request", constants.LogFieldRequestID, hdr.RequestID.String(), constants.LogFieldDestAddr, dstAddr)
71+
klog.InfoS("Start proxy udp request", constants.LogFieldRequestID, hdr.RequestID, constants.LogFieldDestAddr, dstAddr)
7272
udpConn := &xnet.UDPConn{UDPConn: upstreamConn.(*net.UDPConn)}
73-
xnet.ProxyUDP(hdr.RequestID.String(), c, udpConn)
73+
xnet.ProxyUDP(hdr.RequestID, c, udpConn)
7474

7575
default:
76-
klog.InfoS("Unknown protocol", constants.LogFieldRequestID, hdr.RequestID.String(), constants.LogFieldDestAddr, dstAddr, constants.LogFieldProtocol, hdr.Protocol)
76+
klog.InfoS("Unknown protocol", constants.LogFieldRequestID, hdr.RequestID, constants.LogFieldDestAddr, dstAddr, constants.LogFieldProtocol, hdr.Protocol)
7777
}
7878
}
7979

cmd/server/tcp_test.go

+1-2
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,6 @@ import (
99
"testing"
1010
"time"
1111

12-
"github.com/google/uuid"
1312
"github.com/stretchr/testify/require"
1413

1514
"github.com/knight42/krelay/pkg/testutils/tcp"
@@ -41,7 +40,7 @@ func TestHandleTCPConn(t *testing.T) {
4140
return nil, fmt.Errorf("dial: %w", err)
4241
}
4342
hdr := xnet.Header{
44-
RequestID: uuid.New(),
43+
RequestID: xnet.NewRequestID(),
4544
Protocol: xnet.ProtocolTCP,
4645
Port: tsPort,
4746
Addr: xnet.AddrFromHost(tsURL.Hostname()),

go.mod

+1-1
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@ module github.com/knight42/krelay
33
go 1.21
44

55
require (
6-
github.com/google/uuid v1.3.0
76
github.com/spf13/cobra v1.7.0
87
github.com/stretchr/testify v1.8.4
98
k8s.io/api v0.28.0
@@ -29,6 +28,7 @@ require (
2928
github.com/google/go-cmp v0.5.9 // indirect
3029
github.com/google/gofuzz v1.2.0 // indirect
3130
github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510 // indirect
31+
github.com/google/uuid v1.3.0 // indirect
3232
github.com/gregjones/httpcache v0.0.0-20190611155906-901d90724c79 // indirect
3333
github.com/imdario/mergo v0.3.16 // indirect
3434
github.com/inconshreveable/mousetrap v1.1.0 // indirect

pkg/xnet/header.go

+50-17
Original file line numberDiff line numberDiff line change
@@ -4,31 +4,57 @@ import (
44
"encoding/binary"
55
"fmt"
66
"io"
7+
"math/rand"
8+
)
79

8-
"github.com/google/uuid"
10+
const (
11+
lengthAllMandatoryFields = 12 // 1(version) + 2(total length) + 5(request id) + 1(protocol) + 2(port) + 1(addr type)
12+
lengthRequestID = 5
913
)
1014

11-
const totalLenOfOtherFields = 23 // 1(version) + 2(total length) + 16(uuid) + 1(protocol) + 2(port) + 1(addr type)
15+
var letterRunes = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789")
16+
17+
func NewRequestID() string {
18+
b := make([]rune, lengthRequestID)
19+
for i := range b {
20+
b[i] = letterRunes[rand.Intn(len(letterRunes))]
21+
}
22+
return string(b)
23+
}
1224

1325
type Header struct {
1426
Version byte
15-
RequestID uuid.UUID
27+
RequestID string
1628
Protocol byte
1729
Port uint16
1830
Addr Addr
1931
}
2032

2133
func (h *Header) Marshal() []byte {
2234
addrBytes := h.Addr.Marshal()
23-
totalLen := totalLenOfOtherFields + len(addrBytes)
35+
totalLen := lengthAllMandatoryFields + len(addrBytes)
2436
buf := make([]byte, totalLen)
25-
buf[0] = h.Version
26-
binary.BigEndian.PutUint16(buf[1:3], uint16(totalLen))
27-
copy(buf[3:19], h.RequestID[:])
28-
buf[19] = h.Protocol
29-
binary.BigEndian.PutUint16(buf[20:22], h.Port)
30-
buf[22] = h.Addr.typ
31-
copy(buf[23:], addrBytes)
37+
38+
cursor := 0
39+
buf[cursor] = h.Version
40+
cursor++
41+
42+
binary.BigEndian.PutUint16(buf[cursor:cursor+2], uint16(totalLen))
43+
cursor += 2
44+
45+
copy(buf[cursor:cursor+lengthRequestID], h.RequestID)
46+
cursor += lengthRequestID
47+
48+
buf[cursor] = h.Protocol
49+
cursor++
50+
51+
binary.BigEndian.PutUint16(buf[cursor:cursor+2], h.Port)
52+
cursor += 2
53+
54+
buf[cursor] = h.Addr.typ
55+
cursor++
56+
57+
copy(buf[cursor:], addrBytes)
3258
return buf
3359
}
3460

@@ -41,7 +67,7 @@ func (h *Header) FromReader(r io.Reader) error {
4167
h.Version = lengthBuf[0]
4268
// TODO: handle different versions
4369
totalLen := binary.BigEndian.Uint16(lengthBuf[1:])
44-
if totalLen < totalLenOfOtherFields {
70+
if totalLen < lengthAllMandatoryFields {
4571
return fmt.Errorf("body too short: %d", totalLen)
4672
}
4773

@@ -51,12 +77,19 @@ func (h *Header) FromReader(r io.Reader) error {
5177
return fmt.Errorf("read full body: %w", err)
5278
}
5379

54-
reqIDBytes := bodyBuf[:16]
55-
proto := bodyBuf[16]
56-
port := binary.BigEndian.Uint16(bodyBuf[17:19])
57-
h.RequestID, _ = uuid.FromBytes(reqIDBytes)
80+
cursor := 0
81+
reqIDBytes := bodyBuf[cursor : cursor+lengthRequestID]
82+
cursor += lengthRequestID
83+
84+
proto := bodyBuf[cursor]
85+
cursor++
86+
87+
port := binary.BigEndian.Uint16(bodyBuf[cursor : cursor+2])
88+
cursor += 2
89+
90+
h.RequestID = string(reqIDBytes)
5891
h.Protocol = proto
5992
h.Port = port
60-
h.Addr = AddrFromBytes(bodyBuf[19], bodyBuf[20:])
93+
h.Addr = AddrFromBytes(bodyBuf[cursor], bodyBuf[cursor+1:])
6194
return nil
6295
}

pkg/xnet/header_test.go

+11-10
Original file line numberDiff line numberDiff line change
@@ -5,26 +5,27 @@ import (
55
"net"
66
"testing"
77

8-
"github.com/google/uuid"
98
"github.com/stretchr/testify/require"
109
)
1110

11+
var fakeRequestID = "00000"
12+
1213
var headerCases = map[string]struct {
1314
hdr Header
1415
bytes []byte
1516
}{
1617
"host": {
1718
hdr: Header{
1819
Version: 1,
19-
RequestID: uuid.UUID{},
20+
RequestID: fakeRequestID,
2021
Protocol: ProtocolTCP,
2122
Port: 80,
2223
Addr: AddrFromHost("a.com"),
2324
},
2425
bytes: []byte{
2526
1,
26-
0, 0x1c,
27-
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
27+
0, 0x11,
28+
0x30, 0x30, 0x30, 0x30, 0x30,
2829
0,
2930
0, 80,
3031
1,
@@ -34,15 +35,15 @@ var headerCases = map[string]struct {
3435
"ipv4": {
3536
hdr: Header{
3637
Version: 0,
37-
RequestID: uuid.UUID{},
38+
RequestID: fakeRequestID,
3839
Protocol: ProtocolUDP,
3940
Port: 53,
4041
Addr: AddrFromBytes(AddrTypeIP, net.IPv4(192, 168, 1, 1).To4()),
4142
},
4243
bytes: []byte{
4344
0,
44-
0, 0x1b,
45-
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
45+
0, 0x10,
46+
0x30, 0x30, 0x30, 0x30, 0x30,
4647
1,
4748
0, 53,
4849
0,
@@ -52,15 +53,15 @@ var headerCases = map[string]struct {
5253
"ipv6": {
5354
hdr: Header{
5455
Version: 0,
55-
RequestID: uuid.UUID{},
56+
RequestID: fakeRequestID,
5657
Protocol: ProtocolTCP,
5758
Port: 8080,
5859
Addr: AddrFromBytes(AddrTypeIP, net.IPv6loopback),
5960
},
6061
bytes: []byte{
6162
0,
62-
0, 0x27,
63-
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
63+
0, 0x1c,
64+
0x30, 0x30, 0x30, 0x30, 0x30,
6465
0,
6566
0x1f, 0x90,
6667
0,

0 commit comments

Comments
 (0)