Add server-driven time limits

2 files changed, 34 insertions, 5 deletions

diff --git a/app/controllers/proxy.go b/app/controllers/proxy.go
index 9e3ec13..0777c93 100644
--- a/app/controllers/proxy.go
+++ b/app/controllers/proxy.go
@@ -72,6 +72,9 @@ func (h *ProxyHandler) Handle(c echo.Context) error {
                 return c.NoContent(http.StatusUnauthorized)
         }
 
+        // TODO: Set Access-Control-Max-Age header if policy requires time-limited
+        // sessions then terminate the session once that timer expires
+
         wsconn, err := h.Upgrader.Upgrade(c.Response(), c.Request(), nil)
         if err != nil {
                 return err
diff --git a/cmd/client/client.go b/cmd/client/client.go
index 54d7190..b72003d 100644
--- a/cmd/client/client.go
+++ b/cmd/client/client.go
@@ -11,6 +11,8 @@ import (
         "net"
         "net/http"
         "os"
+        "strconv"
+        "time"
 
         "code.crute.us/mcrute/ssh-proxy/app"
         "code.crute.us/mcrute/ssh-proxy/proxy"
@@ -139,18 +141,30 @@ func addCertificateToAgent(conn agent.ExtendedAgent, private any, cert *ssh.Cert
         })
 }
 
-func dialProxyHost(ctx context.Context, oauthToken, proxyHost, host, port string) (io.ReadWriteCloser, error) {
+func dialProxyHost(ctx context.Context, oauthToken, proxyHost, host, port string) (io.ReadWriteCloser, time.Duration, error) {
         addr := fmt.Sprintf("wss://%s/proxy-to/%s/%s", proxyHost, host, port)
 
         hdr := http.Header{}
         hdr.Add("Authorization", fmt.Sprintf("Bearer %s", oauthToken))
 
-        conn, _, err := websocket.DefaultDialer.DialContext(ctx, addr, hdr)
+        conn, resp, err := websocket.DefaultDialer.DialContext(ctx, addr, hdr)
         if err != nil {
-                return nil, err
+                return nil, 0, err
+        }
+
+        // Extract the connection TTL in seconds from the header if present,
+        // otherwise there's no TTL. If we fail to terminate by the TTL the server
+        // will do it for us.
+        var ttl time.Duration
+        if th := resp.Header.Get("Access-Control-Max-Age"); th != "" {
+                tp, err := strconv.Atoi(th)
+                if err != nil {
+                        return nil, 0, err
+                }
+                ttl = time.Duration(tp) * time.Second
         }
 
-        return &proxy.WebsocketReadWriter{W: conn}, nil
+        return &proxy.WebsocketReadWriter{W: conn}, ttl, nil
 }
 
 func fetchOauthToken(ctx context.Context, clientId, proxyHost string) (string, error) {
@@ -218,7 +232,7 @@ func clientMain(cfg app.Config, host, port, username string) {
                 log.Fatalf("Error adding certificate to agent: %s", err)
         }
 
-        ws, err := dialProxyHost(ctx, oauthToken, cfg.ClientHost, host, port)
+        ws, ttl, err := dialProxyHost(ctx, oauthToken, cfg.ClientHost, host, port)
         if err != nil {
                 log.Fatalf("Error dialing proxy host: %s", err)
         }
@@ -229,13 +243,25 @@ func clientMain(cfg app.Config, host, port, username string) {
 
         errc := make(chan error)
 
+        // The server will also force the connection closed but if we do it here we
+        // can give a slightly more friendly error message to the client.
+        if ttl != 0 {
+                log.Printf("Time limited connection, will expire at %s", time.Now().Add(ttl))
+                time.AfterFunc(ttl, func() {
+                        ws.Close()
+                        errc <- fmt.Errorf("Connection time limit has expired")
+                })
+        }
+
         go proxy.CopyWithErrors(os.Stdout, ws, errc)
         go proxy.CopyWithErrors(ws, os.Stdin, errc)
 
         err = <-errc
         if err != nil {
                 log.Printf("Closing client connection: %s", <-errc)
+                os.Exit(1)
         } else {
                 log.Printf("Closing client connection")
+                os.Exit(0)
         }
 }