small api updates (#69)

These updates make it easier to implement and pass custom Session and
Context implementations

No compatibilty breaking, all tests pass

Author: moul

agent.go

@@ -22,8 +22,10 @@ const (
 // client requested agent forwarding
 var contextKeyAgentRequest = &contextKey{"auth-agent-req"}
 
-func setAgentRequested(sess *session) {
-	sess.ctx.SetValue(contextKeyAgentRequest, true)
+// SetAgentRequested sets up the session context so that AgentRequested
+// returns true.
+func SetAgentRequested(ctx Context) {
+	ctx.SetValue(contextKeyAgentRequest, true)
 }
 
 // AgentRequested returns true if the client requested agent forwarding.

context.go

@@ -103,7 +103,7 @@ func newContext(srv *Server) (*sshContext, context.CancelFunc) {
 
 // this is separate from newContext because we will get ConnMetadata
 // at different points so it needs to be applied separately
-func (ctx *sshContext) applyConnMetadata(conn gossh.ConnMetadata) {
+func applyConnMetadata(ctx Context, conn gossh.ConnMetadata) {
 	if ctx.Value(ContextKeySessionID) != nil {
 		return
 	}

server.go

@@ -42,7 +42,7 @@ type Server struct {
 }
 
 // internal for now
-type channelHandler func(srv *Server, conn *gossh.ServerConn, newChan gossh.NewChannel, ctx *sshContext)
+type channelHandler func(srv *Server, conn *gossh.ServerConn, newChan gossh.NewChannel, ctx Context)
 
 func (srv *Server) ensureHostSigner() error {
 	if len(srv.HostSigners) == 0 {
@@ -55,7 +55,7 @@ func (srv *Server) ensureHostSigner() error {
 	return nil
 }
 
-func (srv *Server) config(ctx *sshContext) *gossh.ServerConfig {
+func (srv *Server) config(ctx Context) *gossh.ServerConfig {
 	srv.channelHandlers = map[string]channelHandler{
 		"session":      sessionHandler,
 		"direct-tcpip": directTcpipHandler,
@@ -72,7 +72,7 @@ func (srv *Server) config(ctx *sshContext) *gossh.ServerConfig {
 	}
 	if srv.PasswordHandler != nil {
 		config.PasswordCallback = func(conn gossh.ConnMetadata, password []byte) (*gossh.Permissions, error) {
-			ctx.applyConnMetadata(conn)
+			applyConnMetadata(ctx, conn)
 			if ok := srv.PasswordHandler(ctx, string(password)); !ok {
 				return ctx.Permissions().Permissions, fmt.Errorf("permission denied")
 			}
@@ -81,7 +81,7 @@ func (srv *Server) config(ctx *sshContext) *gossh.ServerConfig {
 	}
 	if srv.PublicKeyHandler != nil {
 		config.PublicKeyCallback = func(conn gossh.ConnMetadata, key gossh.PublicKey) (*gossh.Permissions, error) {
-			ctx.applyConnMetadata(conn)
+			applyConnMetadata(ctx, conn)
 			if ok := srv.PublicKeyHandler(ctx, key); !ok {
 				return ctx.Permissions().Permissions, fmt.Errorf("permission denied")
 			}
@@ -223,7 +223,7 @@ func (srv *Server) handleConn(newConn net.Conn) {
 	defer srv.trackConn(sshConn, false)
 
 	ctx.SetValue(ContextKeyConn, sshConn)
-	ctx.applyConnMetadata(sshConn)
+	applyConnMetadata(ctx, sshConn)
 	go gossh.DiscardRequests(reqs)
 	for ch := range chans {
 		handler, found := srv.channelHandlers[ch.ChannelType()]

session.go

@@ -77,7 +77,7 @@ type Session interface {
 // when there is no signal channel specified
 const maxSigBufSize = 128
 
-func sessionHandler(srv *Server, conn *gossh.ServerConn, newChan gossh.NewChannel, ctx *sshContext) {
+func sessionHandler(srv *Server, conn *gossh.ServerConn, newChan gossh.NewChannel, ctx Context) {
 	ch, reqs, err := newChan.Accept()
 	if err != nil {
 		// TODO: trigger event callback
@@ -105,7 +105,7 @@ type session struct {
 	env     []string
 	ptyCb   PtyCallback
 	cmd     []string
-	ctx     *sshContext
+	ctx     Context
 	sigCh   chan<- Signal
 	sigBuf  []Signal
 }
@@ -142,7 +142,7 @@ func (sess *session) Permissions() Permissions {
 }
 
 func (sess *session) Context() context.Context {
-	return sess.ctx.Context
+	return sess.ctx
 }
 
 func (sess *session) Exit(code int) error {
@@ -278,7 +278,7 @@ func (sess *session) handleRequests(reqs <-chan *gossh.Request) {
 			req.Reply(ok, nil)
 		case agentRequestType:
 			// TODO: option/callback to allow agent forwarding
-			setAgentRequested(sess)
+			SetAgentRequested(sess.ctx)
 			req.Reply(true, nil)
 		default:
 			// TODO: debug log

tcpip.go

@@ -17,7 +17,7 @@ type forwardData struct {
 	OriginatorPort uint32
 }
 
-func directTcpipHandler(srv *Server, conn *gossh.ServerConn, newChan gossh.NewChannel, ctx *sshContext) {
+func directTcpipHandler(srv *Server, conn *gossh.ServerConn, newChan gossh.NewChannel, ctx Context) {
 	d := forwardData{}
 	if err := gossh.Unmarshal(newChan.ExtraData(), &d); err != nil {
 		newChan.Reject(gossh.ConnectionFailed, "error parsing forward data: "+err.Error())