From cc051a47f5826cae1aefcaf80d5edfa51920b8ca Mon Sep 17 00:00:00 2001
From: hackerchai <i@hackerchai.com>
Date: Tue, 19 Dec 2023 23:08:53 +0800
Subject: [PATCH] fix(cli): Add CodeChallenge logic to ensure security

Signed-off-by: hackerchai <i@hackerchai.com>
---
 go.mod                    |  6 +++---
 pkg/auth/client.go        |  2 +-
 pkg/webapp/webapp_flow.go | 24 ++++++++++++++++++------
 3 files changed, 22 insertions(+), 10 deletions(-)

diff --git a/go.mod b/go.mod
index 72b9186..2a4dc8a 100644
--- a/go.mod
+++ b/go.mod
@@ -21,14 +21,16 @@ require (
 )
 
 require (
+	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/google/go-querystring v1.1.0 // indirect
 	github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
+	github.com/stretchr/objx v0.5.0 // indirect
 	gopkg.in/check.v1 v1.0.0-20200227125254-8fa46927fb4f // indirect
 )
 
 require (
 	github.com/benbjohnson/clock v1.1.0 // indirect
-	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/fsnotify/fsnotify v1.6.0 // indirect
 	github.com/go-logr/logr v1.2.4 // indirect
 	github.com/go-task/slim-sprig v0.0.0-20230315185526-52ccab3ef572 // indirect
@@ -48,13 +50,11 @@ require (
 	github.com/mitchellh/mapstructure v1.5.0 // indirect
 	github.com/olekukonko/tablewriter v0.0.5
 	github.com/pelletier/go-toml/v2 v2.0.8 // indirect
-	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/rivo/uniseg v0.2.0 // indirect
 	github.com/spf13/afero v1.9.5 // indirect
 	github.com/spf13/cast v1.5.1 // indirect
 	github.com/spf13/jwalterweatherman v1.1.0 // indirect
 	github.com/spf13/pflag v1.0.5 // indirect
-	github.com/stretchr/objx v0.5.0 // indirect
 	github.com/subosito/gotenv v1.4.2 // indirect
 	go.uber.org/atomic v1.9.0 // indirect
 	go.uber.org/multierr v1.8.0 // indirect
diff --git a/pkg/auth/client.go b/pkg/auth/client.go
index 254aaff..d8fbeac 100644
--- a/pkg/auth/client.go
+++ b/pkg/auth/client.go
@@ -47,7 +47,7 @@ func NewZeaburWebAppOAuthClient() *WebAppClient {
 		ClientID:               ZeaburOAuthCLIClientID,
 		ClientSecret:           ZeaburOAuthCLIClientSecret,
 		RedirectURIWithoutPort: OAuthLocalServerCallbackURL,
-		Scopes:                 []string{"all"},
+		Scopes:                 []string{"project:write", "user:write", "service:write", "environment:write"},
 
 		AuthorizeURL: ZeaburOAuthAuthorizeURL,
 		TokenURL:     ZeaburOAuthTokenURL,
diff --git a/pkg/webapp/webapp_flow.go b/pkg/webapp/webapp_flow.go
index fb0b877..cd1e9ec 100644
--- a/pkg/webapp/webapp_flow.go
+++ b/pkg/webapp/webapp_flow.go
@@ -5,6 +5,8 @@ package webapp
 import (
 	"context"
 	"crypto/rand"
+	"crypto/sha256"
+	"encoding/base64"
 	"encoding/hex"
 	"errors"
 	"fmt"
@@ -17,8 +19,9 @@ import (
 
 // Flow holds the state for the steps of OAuth Web Application flow.
 type Flow struct {
-	server *localServer
-	state  string
+	server        *localServer
+	state         string
+	codeChallenge string
 }
 
 // InitFlow creates a new Flow instance by detecting a locally available port number.
@@ -29,10 +32,12 @@ func InitFlow() (*Flow, error) {
 	}
 
 	state, _ := randomString(20)
+	codeChallenge, _ := randomString(43)
 
 	return &Flow{
-		server: server,
-		state:  state,
+		server:        server,
+		state:         state,
+		codeChallenge: codeChallenge,
 	}, nil
 }
 
@@ -55,7 +60,9 @@ func (flow *Flow) BrowserURL(baseURL string, config oauth2.Config) (string, erro
 	q := url.Values{}
 	q.Set("client_id", config.ClientID)
 	q.Set("redirect_uri", config.RedirectURL)
-	q.Set("scope", strings.Join(config.Scopes, " "))
+	q.Set("code_challenge", genCodeChallengeS256(flow.codeChallenge))
+	q.Set("code_challenge_method", "S256")
+	q.Set("scope", strings.Join(config.Scopes, ","))
 	q.Set("state", flow.state)
 	q.Set("response_type", "code")
 
@@ -79,7 +86,7 @@ func (flow *Flow) Wait(ctx context.Context, config oauth2.Config) (*oauth2.Token
 		return nil, errors.New("state mismatch")
 	}
 
-	token, err := config.Exchange(context.Background(), code.Code)
+	token, err := config.Exchange(context.Background(), code.Code, oauth2.SetAuthURLParam("code_verifier", flow.codeChallenge))
 	if err != nil {
 		return nil, err
 	}
@@ -95,3 +102,8 @@ func randomString(length int) (string, error) {
 	}
 	return hex.EncodeToString(b), nil
 }
+
+func genCodeChallengeS256(s string) string {
+	s256 := sha256.Sum256([]byte(s))
+	return base64.URLEncoding.EncodeToString(s256[:])
+}