Use callback to create net.Listener (#1)

Co-authored-by: Dean Sheather <dean@deansheather.com>

Author: samchouse

ssh.go

@@ -2,6 +2,7 @@ package ssh
 
 import (
 	"crypto/subtle"
+	"errors"
 	"net"
 
 	gossh "golang.org/x/crypto/ssh"
@@ -29,6 +30,9 @@ const (
 // DefaultHandler is the default Handler used by Serve.
 var DefaultHandler Handler
 
+// ErrReject is returned by some callbacks to reject a request.
+var ErrRejected = errors.New("rejected")
+
 // Option is a functional option handler for Server.
 type Option func(*Server) error
 
@@ -69,8 +73,9 @@ type ReversePortForwardingCallback func(ctx Context, bindHost string, bindPort u
 type LocalUnixForwardingCallback func(ctx Context, socketPath string) bool
 
 // ReverseUnixForwardingCallback is a hook for allowing reverse unix forwarding
-// (streamlocal-forward@openssh.com).
-type ReverseUnixForwardingCallback func(ctx Context, socketPath string) bool
+// (streamlocal-forward@openssh.com). Returning ErrRejected will reject the
+// request.
+type ReverseUnixForwardingCallback func(ctx Context, socketPath string) (net.Listener, error)
 
 // ServerConfigCallback is a hook for creating custom default server configs
 type ServerConfigCallback func(ctx Context) *gossh.ServerConfig

streamlocal.go

@@ -110,7 +110,7 @@ func (h *ForwardedUnixHandler) HandleSSHRequest(ctx Context, srv *Server, req *g
 			return false, nil
 		}
 
-		if srv.ReverseUnixForwardingCallback == nil || !srv.ReverseUnixForwardingCallback(ctx, reqPayload.SocketPath) {
+		if srv.ReverseUnixForwardingCallback == nil {
 			return false, []byte("unix forwarding is disabled")
 		}
 
@@ -123,26 +123,11 @@ func (h *ForwardedUnixHandler) HandleSSHRequest(ctx Context, srv *Server, req *g
 			return false, nil
 		}
 
-		// Create socket parent dir if not exists.
-		parentDir := filepath.Dir(addr)
-		err = os.MkdirAll(parentDir, 0700)
-		if err != nil {
-			// TODO: log mkdir failure
-			return false, nil
-		}
-
-		// Remove existing socket if it exists. We do not use os.Remove() here
-		// so that directories are kept. Note that it's possible that we will
-		// overwrite a regular file here. Both of these behaviors match OpenSSH,
-		// however, which is why we unlink.
-		err = unlink(addr)
-		if err != nil && !errors.Is(err, fs.ErrNotExist) {
-			// TODO: log unlink failure
-			return false, nil
-		}
-
-		ln, err := net.Listen("unix", addr)
+		ln, err := srv.ReverseUnixForwardingCallback(ctx, addr)
 		if err != nil {
+			if errors.Is(err, ErrRejected) {
+				return false, []byte("unix forwarding is disabled")
+			}
 			// TODO: log unix listen failure
 			return false, nil
 		}
@@ -227,3 +212,31 @@ func unlink(path string) error {
 		}
 	}
 }
+
+// SimpleUnixReverseForwardingCallback provides a basic implementation for
+// ReverseUnixForwardingCallback. The parent directory will be created (with
+// os.MkdirAll), and existing files with the same name will be removed.
+func SimpleUnixReverseForwardingCallback(_ Context, socketPath string) (net.Listener, error) {
+	// Create socket parent dir if not exists.
+	parentDir := filepath.Dir(socketPath)
+	err := os.MkdirAll(parentDir, 0700)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create parent directory %q for socket %q: %w", parentDir, socketPath, err)
+	}
+
+	// Remove existing socket if it exists. We do not use os.Remove() here
+	// so that directories are kept. Note that it's possible that we will
+	// overwrite a regular file here. Both of these behaviors match OpenSSH,
+	// however, which is why we unlink.
+	err = unlink(socketPath)
+	if err != nil && !errors.Is(err, fs.ErrNotExist) {
+		return nil, fmt.Errorf("failed to remove existing file in socket path %q: %w", socketPath, err)
+	}
+
+	ln, err := net.Listen("unix", socketPath)
+	if err != nil {
+		return nil, fmt.Errorf("failed to listen on unix socket %q: %w", socketPath, err)
+	}
+
+	return ln, err
+}

streamlocal_test.go

@@ -127,11 +127,11 @@ func TestReverseUnixForwardingWorks(t *testing.T) {
 
 	_, client, cleanup := newTestSession(t, &Server{
 		Handler: func(s Session) {},
-		ReverseUnixForwardingCallback: func(ctx Context, socketPath string) bool {
+		ReverseUnixForwardingCallback: func(ctx Context, socketPath string) (net.Listener, error) {
 			if socketPath != remoteSocketPath {
 				panic("unexpected socket path: " + socketPath)
 			}
-			return true
+			return SimpleUnixReverseForwardingCallback(ctx, socketPath)
 		},
 	}, nil)
 	defer cleanup()
@@ -182,12 +182,12 @@ func TestReverseUnixForwardingRespectsCallback(t *testing.T) {
 	var called int64
 	_, client, cleanup := newTestSession(t, &Server{
 		Handler: func(s Session) {},
-		ReverseUnixForwardingCallback: func(ctx Context, socketPath string) bool {
+		ReverseUnixForwardingCallback: func(ctx Context, socketPath string) (net.Listener, error) {
 			atomic.AddInt64(&called, 1)
 			if socketPath != remoteSocketPath {
 				panic("unexpected socket path: " + socketPath)
 			}
-			return false
+			return nil, ErrRejected
 		},
 	}, nil)
 	defer cleanup()