From 6a9001bdb1603d1a4eca8eeebd4adbc1e2bc053f Mon Sep 17 00:00:00 2001
From: ben-toogood <bentoogood@gmail.com>
Date: Wed, 4 Mar 2020 09:54:52 +0000
Subject: [PATCH] Set auth account in context (#1293)

---
 auth/auth.go            | 42 +++++++++++++++++++++++++++++++++++++++++
 util/wrapper/wrapper.go | 32 ++++++++++++++++++++++---------
 2 files changed, 65 insertions(+), 9 deletions(-)

diff --git a/auth/auth.go b/auth/auth.go
index f98b2a67..06ef3b5d 100644
--- a/auth/auth.go
+++ b/auth/auth.go
@@ -2,7 +2,11 @@
 package auth
 
 import (
+	"context"
+	"encoding/json"
 	"time"
+
+	"github.com/micro/go-micro/v2/metadata"
 )
 
 // Auth providers authentication and authorization
@@ -53,3 +57,41 @@ type Account struct {
 	// Any other associated metadata
 	Metadata map[string]string `json:"metadata"`
 }
+
+const (
+	// MetadataKey is the key used when storing the account
+	// in metadata
+	MetadataKey = "auth-account"
+)
+
+// AccountFromContext gets the account from the context, which
+// is set by the auth wrapper at the start of a call. If the account
+// is not set, a nil account will be returned. The error is only returned
+// when there was a problem retrieving an account
+func AccountFromContext(ctx context.Context) (*Account, error) {
+	str, ok := metadata.Get(ctx, MetadataKey)
+	// there was no account set
+	if !ok {
+		return nil, nil
+	}
+
+	var acc *Account
+	// metadata is stored as a string, so unmarshal to an account
+	if err := json.Unmarshal([]byte(str), &acc); err != nil {
+		return nil, err
+	}
+
+	return acc, nil
+}
+
+// ContextWithAccount sets the account in the context
+func ContextWithAccount(ctx context.Context, account *Account) (context.Context, error) {
+	// metadata is stored as a string, so marshal to bytes
+	bytes, err := json.Marshal(account)
+	if err != nil {
+		return ctx, err
+	}
+
+	// generate a new context with the MetadataKey set
+	return metadata.Set(ctx, MetadataKey, string(bytes)), nil
+}
diff --git a/util/wrapper/wrapper.go b/util/wrapper/wrapper.go
index b83aeca0..80dc702b 100644
--- a/util/wrapper/wrapper.go
+++ b/util/wrapper/wrapper.go
@@ -164,13 +164,6 @@ func AuthHandler(fn func() auth.Auth) server.HandlerWrapper {
 				return h(ctx, req, rsp)
 			}
 
-			// Exclude any user excluded endpoints
-			for _, e := range a.Options().Exclude {
-				if e == req.Endpoint() {
-					return h(ctx, req, rsp)
-				}
-			}
-
 			// Extract the token if present. Note: if noop is being used
 			// then the token can be blank without erroring
 			var token string
@@ -184,10 +177,31 @@ func AuthHandler(fn func() auth.Auth) server.HandlerWrapper {
 			}
 
 			// Verify the token
-			if _, err := a.Verify(token); err != nil {
-				return errors.Unauthorized("go.micro.auth", err.Error())
+			account, authErr := a.Verify(token)
+
+			// If there is an account, set it in the context
+			if authErr == nil {
+				var err error
+				ctx, err = auth.ContextWithAccount(ctx, account)
+
+				if err != nil {
+					return err
+				}
 			}
 
+			// Return if the user disabled auth on this endpoint
+			for _, e := range a.Options().Exclude {
+				if e == req.Endpoint() {
+					return h(ctx, req, rsp)
+				}
+			}
+
+			// If the authErr is set, prevent the user from calling the endpoint
+			if authErr != nil {
+				return errors.Unauthorized("go.micro.auth", authErr.Error())
+			}
+
+			// The user is authorised, allow the call
 			return h(ctx, req, rsp)
 		}
 	}