From 9180d4bd459e3f7efe04c7e582dad2b8df7fcb8e Mon Sep 17 00:00:00 2001
From: YR Chen <yrchen@mail.ustc.edu.cn>
Date: Sat, 24 Aug 2024 22:21:36 +0800
Subject: [PATCH] Change API response `proxy_protocol` to string

---
 auth.go        |  2 +-
 legacy_auth.go | 15 +++++++++------
 sshmux.go      | 17 +++++++++++++----
 sshmux_test.go |  3 ++-
 4 files changed, 25 insertions(+), 12 deletions(-)

diff --git a/auth.go b/auth.go
index c47d9e5..26f77d8 100644
--- a/auth.go
+++ b/auth.go
@@ -46,7 +46,7 @@ type AuthUpstream struct {
 	PrivateKey    string  `json:"private_key,omitempty"`
 	Certificate   string  `json:"certificate,omitempty"`
 	Password      *string `json:"password,omitempty"`
-	ProxyProtocol byte    `json:"proxy_protocol,omitempty"`
+	ProxyProtocol *string `json:"proxy_protocol,omitempty"`
 }
 
 type Authenticator interface {
diff --git a/legacy_auth.go b/legacy_auth.go
index a7cb61e..a16b3fd 100644
--- a/legacy_auth.go
+++ b/legacy_auth.go
@@ -128,17 +128,20 @@ func (auth *LegacyAuthenticator) Auth(request AuthRequest, username string) (int
 			return 500, nil, err
 		}
 		auth_upstream := AuthUpstream{
-			Host:          address.Addr().String(),
-			Port:          address.Port(),
-			PrivateKey:    upstream.PrivateKey,
-			Certificate:   upstream.Certificate,
-			Password:      upstream.Password,
-			ProxyProtocol: upstream.ProxyProtocol,
+			Host:        address.Addr().String(),
+			Port:        address.Port(),
+			PrivateKey:  upstream.PrivateKey,
+			Certificate: upstream.Certificate,
+			Password:    upstream.Password,
 		}
 		unix_password, has_unix_password := request.Payload["unix_password"]
 		if has_unix_password {
 			auth_upstream.Password = &unix_password
 		}
+		if upstream.ProxyProtocol > 0 {
+			proxyProtocol := fmt.Sprintf("v%d", upstream.ProxyProtocol)
+			auth_upstream.ProxyProtocol = &proxyProtocol
+		}
 		resp := AuthResponse{Upstream: &auth_upstream}
 		return 200, &resp, nil
 	}
diff --git a/sshmux.go b/sshmux.go
index 981847e..fde9cd7 100644
--- a/sshmux.go
+++ b/sshmux.go
@@ -219,10 +219,19 @@ auth_requests:
 					upstreamResp.Port = 22
 				}
 				upstream = &UpstreamInformation{
-					Host:          netip.AddrPortFrom(host, upstreamResp.Port).String(),
-					Signer:        parsePrivateKey(upstreamResp.PrivateKey, upstreamResp.Certificate),
-					Password:      upstreamResp.Password,
-					ProxyProtocol: upstreamResp.ProxyProtocol,
+					Host:     netip.AddrPortFrom(host, upstreamResp.Port).String(),
+					Signer:   parsePrivateKey(upstreamResp.PrivateKey, upstreamResp.Certificate),
+					Password: upstreamResp.Password,
+				}
+				if upstreamResp.ProxyProtocol != nil {
+					switch *upstreamResp.ProxyProtocol {
+					case "v1":
+						upstream.ProxyProtocol = 1
+					case "v2":
+						upstream.ProxyProtocol = 2
+					default:
+						return fmt.Errorf("unknown PROXY protocol version: %s", *upstreamResp.ProxyProtocol)
+					}
 				}
 				break auth_requests
 			case 401:
diff --git a/sshmux_test.go b/sshmux_test.go
index 622d16a..14baa35 100644
--- a/sshmux_test.go
+++ b/sshmux_test.go
@@ -82,9 +82,10 @@ func initHttp(sshPrivateKey []byte) {
 			PrivateKey: string(sshPrivateKey),
 		}
 		if enableProxy {
+			proxyProtocol := "v2"
 			upstream.Host = sshdProxiedAddr.IP.String()
 			upstream.Port = uint16(sshdProxiedAddr.Port)
-			upstream.ProxyProtocol = 2
+			upstream.ProxyProtocol = &proxyProtocol
 		} else {
 			upstream.Host = sshdServerAddr.IP.String()
 			upstream.Port = uint16(sshdServerAddr.Port)