Merge pull request digitalocean#421 from andrewsomething/APICLI-327/t…

…oken-validation Add some basic input cleaning to NewFromToken
dmartzol · Dec 2, 2020 · 05dc4d9 · 05dc4d9
2 parents a11a281 + fcbbb1e
commit 05dc4d9
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 2 deletions.
diff --git a/godo.go b/godo.go
@@ -11,6 +11,7 @@ import (
 	"net/url"
 	"reflect"
 	"strconv"
+	"strings"
 	"sync"
 	"time"
 
@@ -168,8 +169,9 @@ func addOptions(s string, opt interface{}) (string, error) {
 // NewFromToken returns a new DigitalOcean API client with the given API
 // token.
 func NewFromToken(token string) *Client {
+	cleanToken := strings.Trim(strings.TrimSpace(token), "'")
 	ctx := context.Background()
-	ts := oauth2.StaticTokenSource(&oauth2.Token{AccessToken: token})
+	ts := oauth2.StaticTokenSource(&oauth2.Token{AccessToken: cleanToken})
 	return NewClient(oauth2.NewClient(ctx, ts))
 }
 

diff --git a/godo_test.go b/godo_test.go
@@ -126,10 +126,38 @@ func TestNewClient(t *testing.T) {
 }
 
 func TestNewFromToken(t *testing.T) {
-	c := NewFromToken("my-token")
+	c := NewFromToken("myToken")
 	testClientDefaults(t, c)
 }
 
+func TestNewFromToken_cleaned(t *testing.T) {
+	testTokens := []string{"myToken ", " myToken", " myToken ", "'myToken'", " 'myToken' "}
+	expected := "Bearer myToken"
+
+	setup()
+	defer teardown()
+
+	mux.HandleFunc("/foo", func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+	})
+
+	for _, tt := range testTokens {
+		t.Run(tt, func(t *testing.T) {
+			c := NewFromToken(tt)
+			req, _ := c.NewRequest(ctx, http.MethodGet, server.URL+"/foo", nil)
+			resp, err := c.Do(ctx, req, nil)
+			if err != nil {
+				t.Fatalf("Do(): %v", err)
+			}
+
+			authHeader := resp.Request.Header.Get("Authorization")
+			if authHeader != expected {
+				t.Errorf("Authorization header = %v, expected %v", authHeader, expected)
+			}
+		})
+	}
+}
+
 func TestNew(t *testing.T) {
 	c, err := New(nil)