Merge pull request #137 from uselagoon/token-grant

feat: add support in ssh-token for "grant" command
uselagoon · Nov 24, 2022 · f3d5f99 · f3d5f99
2 parents bb28a7f + 250f562
commit f3d5f99
Show file tree

Hide file tree

Showing 2 changed files with 98 additions and 25 deletions.
diff --git a/internal/keycloak/useraccesstoken.go b/internal/keycloak/useraccesstoken.go
@@ -2,6 +2,7 @@ package keycloak
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"net/http"
 	"path"
@@ -12,13 +13,10 @@ import (
 	"golang.org/x/oauth2"
 )
 
-// UserAccessToken queries Keycloak given the user UUID, and returns an access
-// token. Authorized party for this token is auth-server. Authorization is done
-// by the Lagoon API.
-func (c *Client) UserAccessToken(ctx context.Context,
-	userUUID *uuid.UUID) (string, error) {
+func (c *Client) getUserToken(ctx context.Context,
+	userUUID *uuid.UUID) (*oauth2.Token, error) {
 	// set up tracing
-	ctx, span := otel.Tracer(pkgName).Start(ctx, "UserAccessToken")
+	ctx, span := otel.Tracer(pkgName).Start(ctx, "getUserToken")
 	defer span.End()
 	// get user token
 	tokenURL := *c.baseURL
@@ -41,12 +39,49 @@ func (c *Client) UserAccessToken(ctx context.Context,
 		// https://www.keycloak.org/docs/latest/securing_apps/#_token-exchange
 		oauth2.SetAuthURLParam("requested_subject", userUUID.String()))
 	if err != nil {
-		return "", fmt.Errorf("couldn't get user token: %v", err)
+		return nil, fmt.Errorf("couldn't get user token: %v", err)
 	}
 	// parse and extract verified attributes
 	_, err = c.validateTokenClaims(userToken)
 	if err != nil {
-		return "", fmt.Errorf("couldn't validate token claims: %v", err)
+		return nil, fmt.Errorf("couldn't validate token claims: %v", err)
+	}
+	return userToken, nil
+}
+
+// UserAccessTokenResponse queries Keycloak given the user UUID, and returns an
+// access token response containing both access_token and refresh_token.
+// Authorized party for these tokens is auth-server. Authorization is done by
+// the Lagoon API.
+func (c *Client) UserAccessTokenResponse(ctx context.Context,
+	userUUID *uuid.UUID) (string, error) {
+	// set up tracing
+	ctx, span := otel.Tracer(pkgName).Start(ctx, "UserAccessToken")
+	defer span.End()
+	// get user token
+	userToken, err := c.getUserToken(ctx, userUUID)
+	if err != nil {
+		return "", fmt.Errorf("couldn't get user token: %v", err)
+	}
+	data, err := json.Marshal(userToken)
+	if err != nil {
+		return "", fmt.Errorf("couldn't marshal user token: %v", err)
+	}
+	return string(data), nil
+}
+
+// UserAccessToken queries Keycloak given the user UUID, and returns an access
+// token. Authorized party for this token is auth-server. Authorization is done
+// by the Lagoon API.
+func (c *Client) UserAccessToken(ctx context.Context,
+	userUUID *uuid.UUID) (string, error) {
+	// set up tracing
+	ctx, span := otel.Tracer(pkgName).Start(ctx, "UserAccessToken")
+	defer span.End()
+	// get user token
+	userToken, err := c.getUserToken(ctx, userUUID)
+	if err != nil {
+		return "", fmt.Errorf("couldn't get user token: %v", err)
 	}
 	return userToken.AccessToken, nil
 }
diff --git a/internal/sshtoken/sessionhandler.go b/internal/sshtoken/sessionhandler.go
@@ -13,6 +13,7 @@ import (
 
 // KeycloakService provides methods for querying the Keycloak API.
 type KeycloakService interface {
+	UserAccessTokenResponse(context.Context, *uuid.UUID) (string, error)
 	UserAccessToken(context.Context, *uuid.UUID) (string, error)
 }
 
@@ -57,48 +58,85 @@ func sessionHandler(log *zap.Logger, k KeycloakService) ssh.Handler {
 			}
 			return
 		}
-		// validate the command. right now we only support "token".
+		// valid commands:
+		// - grant: returns a full access token response as per
+		//   https://www.rfc-editor.org/rfc/rfc6749#section-4.1.4
+		// - token: returns a bare access token (the contents of the access_token
+		//   field inside a full token access token response)
 		cmd := s.Command()
-		if len(cmd) != 1 || cmd[0] != "token" {
-			log.Debug("invalid command",
+		if len(cmd) != 1 {
+			log.Debug("too many arguments",
 				zap.Strings("command", cmd),
 				zap.String("sessionID", sid))
 			_, err := fmt.Fprintf(s.Stderr(),
-				"invalid command: only \"token\" is supported. SID: %s\n", sid)
+				"invalid command: only a single argument is supported. SID: %s\n", sid)
 			if err != nil {
 				log.Debug("couldn't write error message to session stream",
 					zap.String("sessionID", sid),
 					zap.Error(err))
 			}
 			return
 		}
-		// get the user access token from keycloak
-		accessToken, err := k.UserAccessToken(s.Context(), uid)
-		if err != nil {
-			log.Warn("couldn't get user access token",
-				zap.String("sessionID", sid),
-				zap.String("userUUID", uid.String()),
-				zap.Error(err))
-			_, err = fmt.Fprintf(s.Stderr(),
-				"internal error. SID: %s\n", sid)
+		// get response
+		var response string
+		var err error
+		switch cmd[0] {
+		case "grant":
+			response, err = k.UserAccessTokenResponse(s.Context(), uid)
+			if err != nil {
+				log.Warn("couldn't get user access token response",
+					zap.String("sessionID", sid),
+					zap.String("userUUID", uid.String()),
+					zap.Error(err))
+				_, err = fmt.Fprintf(s.Stderr(),
+					"internal error. SID: %s\n", sid)
+				if err != nil {
+					log.Debug("couldn't write error message to session stream",
+						zap.String("sessionID", sid),
+						zap.Error(err))
+				}
+				return
+			}
+		case "token":
+			response, err = k.UserAccessToken(s.Context(), uid)
+			if err != nil {
+				log.Warn("couldn't get user access token",
+					zap.String("sessionID", sid),
+					zap.String("userUUID", uid.String()),
+					zap.Error(err))
+				_, err = fmt.Fprintf(s.Stderr(),
+					"internal error. SID: %s\n", sid)
+				if err != nil {
+					log.Debug("couldn't write error message to session stream",
+						zap.String("sessionID", sid),
+						zap.Error(err))
+				}
+				return
+			}
+		default:
+			log.Debug("invalid command",
+				zap.Strings("command", cmd),
+				zap.String("sessionID", sid))
+			_, err := fmt.Fprintf(s.Stderr(),
+				"invalid command: only \"grant\" and \"token\" are supported. SID: %s\n", sid)
 			if err != nil {
 				log.Debug("couldn't write error message to session stream",
 					zap.String("sessionID", sid),
 					zap.Error(err))
 			}
 			return
 		}
-		// send token response
-		_, err = fmt.Fprintf(s, "%s\n", accessToken)
+		// send response
+		_, err = fmt.Fprintf(s, "%s\n", response)
 		if err != nil {
-			log.Debug("couldn't write token to session stream",
+			log.Debug("couldn't write response to session stream",
 				zap.String("sessionID", sid),
 				zap.String("userUUID", uid.String()),
 				zap.Error(err))
 			return
 		}
 		tokensGeneratedTotal.Inc()
-		log.Info("generated token for user",
+		log.Info("generated access token for user",
 			zap.String("sessionID", sid),
 			zap.String("userUUID", uid.String()))
 	}