From 76260c3b48b08879e631a3e410b5db4e473653e0 Mon Sep 17 00:00:00 2001
From: Stan Hu <stanhu@gmail.com>
Date: Mon, 26 Jul 2021 15:54:21 -0700
Subject: Set a 90-second timeout on proxy headers

---
 internal/sshd/sshd.go | 21 ++++++++++++---------
 1 file changed, 12 insertions(+), 9 deletions(-)

diff --git a/internal/sshd/sshd.go b/internal/sshd/sshd.go
index d3b5ec1..8b49712 100644
--- a/internal/sshd/sshd.go
+++ b/internal/sshd/sshd.go
@@ -7,10 +7,10 @@ import (
 	"fmt"
 	"io/ioutil"
 	"net"
+	"net/http"
 	"strconv"
-	"time"
 	"sync"
-	"net/http"
+	"time"
 
 	"github.com/pires/go-proxyproto"
 	"golang.org/x/crypto/ssh"
@@ -18,25 +18,26 @@ import (
 	"gitlab.com/gitlab-org/gitlab-shell/internal/config"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/gitlabnet/authorizedkeys"
 
-	"gitlab.com/gitlab-org/labkit/log"
 	"gitlab.com/gitlab-org/labkit/correlation"
+	"gitlab.com/gitlab-org/labkit/log"
 )
 
 type status int
 
-const(
+const (
 	StatusStarting status = iota
 	StatusReady
 	StatusOnShutdown
 	StatusClosed
+	ProxyHeaderTimeout = 90 * time.Second
 )
 
 type Server struct {
 	Config *config.Config
 
-	status status
+	status   status
 	statusMu sync.Mutex
-	wg sync.WaitGroup
+	wg       sync.WaitGroup
 	listener net.Listener
 }
 
@@ -71,7 +72,7 @@ func (s *Server) MonitoringServeMux() *http.ServeMux {
 	})
 
 	mux.HandleFunc(s.Config.Server.LivenessProbe, func(w http.ResponseWriter, r *http.Request) {
-			w.WriteHeader(http.StatusOK)
+		w.WriteHeader(http.StatusOK)
 	})
 
 	return mux
@@ -84,7 +85,10 @@ func (s *Server) listen() error {
 	}
 
 	if s.Config.Server.ProxyProtocol {
-		sshListener = &proxyproto.Listener{Listener: sshListener}
+		sshListener = &proxyproto.Listener{
+			Listener:          sshListener,
+			ReadHeaderTimeout: ProxyHeaderTimeout,
+		}
 
 		log.Info("Proxy protocol is enabled")
 	}
@@ -191,7 +195,6 @@ func (s *Server) initConfig(ctx context.Context) (*ssh.ServerConfig, error) {
 	return sshCfg, nil
 }
 
-
 func (s *Server) handleConn(ctx context.Context, sshCfg *ssh.ServerConfig, nconn net.Conn) {
 	remoteAddr := nconn.RemoteAddr().String()
 
-- 
cgit v1.2.1