feat: make server acknowledge request (#23)

Author: knight42

cmd/client/main.go

@@ -224,7 +224,7 @@ service, ip and hostname rather than only pods.`,
 	cf.AddFlags(flags)
 	flags.BoolVarP(&printVersion, "version", "V", false, "Print version info and exit.")
 	flags.StringVar(&o.address, "address", "127.0.0.1", "Address to listen on. Only accepts IP addresses as a value.")
-	flags.StringVar(&o.serverImage, "server.image", constants.ServerImage, "The krelay-server image to use.")
+	flags.StringVar(&o.serverImage, "server.image", "ghcr.io/knight42/krelay-server:v0.0.2", "The krelay-server image to use.")
 	flags.StringVar(&o.serverNamespace, "server.namespace", metav1.NamespaceDefault, "The namespace in which krelay-server is located.")
 
 	// I do not want these flags to show up in --help.

cmd/client/tcp.go

@@ -43,6 +43,17 @@ func handleTCPConn(clientConn net.Conn, serverConn httpstream.Connection, dstAdd
 		return
 	}
 
+	var ack xnet.Acknowledgement
+	err = ack.FromReader(dataStream)
+	if err != nil {
+		klog.ErrorS(err, "Fail to receive ack", kvs...)
+		return
+	}
+	if ack.Code != xnet.AckCodeOK {
+		klog.ErrorS(ack.Code, "Fail to connect", kvs...)
+		return
+	}
+
 	localError := make(chan struct{})
 	remoteDone := make(chan struct{})
 

cmd/client/udp.go

@@ -43,6 +43,17 @@ func handleUDPConn(clientConn net.PacketConn, cliAddr net.Addr, dataCh chan []by
 		return
 	}
 
+	var ack xnet.Acknowledgement
+	err = ack.FromReader(dataStream)
+	if err != nil {
+		klog.ErrorS(err, "Fail to receive ack", kvs...)
+		return
+	}
+	if ack.Code != xnet.AckCodeOK {
+		klog.ErrorS(ack.Code, "Fail to connect", kvs...)
+		return
+	}
+
 	upClosed := make(chan struct{})
 	go func() {
 		var (

cmd/server/main.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"context"
+	"errors"
 	"flag"
 	"fmt"
 	"net"
@@ -40,6 +41,31 @@ func (o *options) run(ctx context.Context) error {
 	}
 }
 
+func writeACK(c net.Conn, ack xnet.Acknowledgement) error {
+	data := ack.Marshal()
+	_, err := c.Write(data)
+	return err
+}
+
+func ackCodeFromErr(err error) xnet.AckCode {
+	var dnsErr *net.DNSError
+	if errors.As(err, &dnsErr) {
+		if dnsErr.IsNotFound {
+			return xnet.AckCodeNoSuchHost
+		}
+		if dnsErr.IsTimeout {
+			return xnet.AckCodeResolveTimeout
+		}
+	}
+
+	var opErr *net.OpError
+	if errors.As(err, &opErr) && opErr.Timeout() {
+		return xnet.AckCodeConnectTimeout
+	}
+
+	return xnet.AckCodeUnknownError
+}
+
 func handleConn(ctx context.Context, c *net.TCPConn, dialer *net.Dialer) {
 	defer c.Close()
 
@@ -57,6 +83,16 @@ func handleConn(ctx context.Context, c *net.TCPConn, dialer *net.Dialer) {
 		upstreamConn, err := dialer.DialContext(ctx, constants.ProtocolTCP, dstAddr)
 		if err != nil {
 			klog.ErrorS(err, "Fail to create tcp connection", constants.LogFieldRequestID, hdr.RequestID, constants.LogFieldDestAddr, dstAddr)
+			_ = writeACK(c, xnet.Acknowledgement{
+				Code: ackCodeFromErr(err),
+			})
+			return
+		}
+		err = writeACK(c, xnet.Acknowledgement{
+			Code: xnet.AckCodeOK,
+		})
+		if err != nil {
+			klog.ErrorS(err, "Fail to write ack", constants.LogFieldRequestID, hdr.RequestID)
 			return
 		}
 		klog.InfoS("Start proxy tcp request", constants.LogFieldRequestID, hdr.RequestID, constants.LogFieldDestAddr, dstAddr)
@@ -66,6 +102,16 @@ func handleConn(ctx context.Context, c *net.TCPConn, dialer *net.Dialer) {
 		upstreamConn, err := dialer.DialContext(ctx, constants.ProtocolUDP, dstAddr)
 		if err != nil {
 			klog.ErrorS(err, "Fail to create udp connection", constants.LogFieldRequestID, hdr.RequestID, constants.LogFieldDestAddr, dstAddr)
+			_ = writeACK(c, xnet.Acknowledgement{
+				Code: ackCodeFromErr(err),
+			})
+			return
+		}
+		err = writeACK(c, xnet.Acknowledgement{
+			Code: xnet.AckCodeOK,
+		})
+		if err != nil {
+			klog.ErrorS(err, "Fail to write ack", constants.LogFieldRequestID, hdr.RequestID)
 			return
 		}
 		klog.InfoS("Start proxy udp request", constants.LogFieldRequestID, hdr.RequestID, constants.LogFieldDestAddr, dstAddr)

cmd/server/tcp_test.go

@@ -49,6 +49,14 @@ func TestHandleTCPConn(t *testing.T) {
 		if err != nil {
 			return nil, fmt.Errorf("write header: %w", err)
 		}
+		var ack xnet.Acknowledgement
+		err = ack.FromReader(c)
+		if err != nil {
+			return nil, fmt.Errorf("read ack: %w", err)
+		}
+		if ack.Code != xnet.AckCodeOK {
+			return nil, fmt.Errorf("ack: %s", ack.Code.Error())
+		}
 		return c, nil
 	}
 

pkg/constants/constants.go

@@ -24,5 +24,3 @@ const (
 	ProtocolTCP = "tcp"
 	ProtocolUDP = "udp"
 )
-
-const ServerImage = "ghcr.io/knight42/krelay-server:v0.0.1"

pkg/xnet/ack.go

@@ -0,0 +1,48 @@
+package xnet
+
+import (
+	"fmt"
+	"io"
+)
+
+type AckCode uint8
+
+const (
+	AckCodeOK = iota + 1
+	AckCodeUnknownError
+	AckCodeNoSuchHost
+	AckCodeResolveTimeout
+	AckCodeConnectTimeout
+)
+
+func (c AckCode) Error() string {
+	switch c {
+	case AckCodeUnknownError:
+		return "Unknown error"
+	case AckCodeNoSuchHost:
+		return "No such host"
+	case AckCodeResolveTimeout:
+		return "Resolve timeout"
+	case AckCodeConnectTimeout:
+		return "Connect timeout"
+	}
+	return "Unknown Code"
+}
+
+type Acknowledgement struct {
+	Code AckCode
+}
+
+func (a *Acknowledgement) Marshal() []byte {
+	return []byte{byte(a.Code)}
+}
+
+func (a *Acknowledgement) FromReader(r io.Reader) error {
+	var buf [1]byte
+	_, err := r.Read(buf[:])
+	if err != nil {
+		return fmt.Errorf("read ack: %w", err)
+	}
+	a.Code = AckCode(buf[0])
+	return nil
+}