1 files changed, 22 insertions, 7 deletions
diff --git a/internal/sshd/sshd.go b/internal/sshd/sshd.go
index b08b386..d20286a 100644
--- a/internal/sshd/sshd.go
+++ b/internal/sshd/sshd.go
@@ -9,7 +9,7 @@ import (
 	"sync"
 	"time"
 
-	"github.com/pires/go-proxyproto"
+	proxyproto "github.com/pires/go-proxyproto"
 	"golang.org/x/crypto/ssh"
 
 	"gitlab.com/gitlab-org/gitlab-shell/v14/client"
@@ -95,9 +95,14 @@ func (s *Server) listen(ctx context.Context) error {
 	}
 
 	if s.Config.Server.ProxyProtocol {
+		policy, err := s.proxyPolicy()
+		if err != nil {
+			return fmt.Errorf("invalid policy configuration: %w", err)
+		}
+
 		sshListener = &proxyproto.Listener{
 			Listener:          sshListener,
-			Policy:            s.requirePolicy,
+			Policy:            policy,
 			ReadHeaderTimeout: time.Duration(s.Config.Server.ProxyHeaderTimeout),
 		}
 
@@ -200,17 +205,27 @@ func (s *Server) handleConn(ctx context.Context, nconn net.Conn) {
 	})
 }
 
-func (s *Server) requirePolicy(_ net.Addr) (proxyproto.Policy, error) {
+func (s *Server) proxyPolicy() (proxyproto.PolicyFunc, error) {
+	if len(s.Config.Server.ProxyAllowed) > 0 {
+		return proxyproto.StrictWhiteListPolicy(s.Config.Server.ProxyAllowed)
+	}
+
 	// Set the Policy value based on config
 	// Values are taken from https://github.com/pires/go-proxyproto/blob/195fedcfbfc1be163f3a0d507fac1709e9d81fed/policy.go#L20
 	switch strings.ToLower(s.Config.Server.ProxyPolicy) {
 	case "require":
-		return proxyproto.REQUIRE, nil
+		return staticProxyPolicy(proxyproto.REQUIRE), nil
 	case "ignore":
-		return proxyproto.IGNORE, nil
+		return staticProxyPolicy(proxyproto.IGNORE), nil
 	case "reject":
-		return proxyproto.REJECT, nil
+		return staticProxyPolicy(proxyproto.REJECT), nil
 	default:
-		return proxyproto.USE, nil
+		return staticProxyPolicy(proxyproto.USE), nil
+	}
+}
+
+func staticProxyPolicy(policy proxyproto.Policy) proxyproto.PolicyFunc {
+	return func(_ net.Addr) (proxyproto.Policy, error) {
+		return policy, nil
 	}
 }