|
| 1 | +// Copyright 2024 PingCAP, Inc. |
| 2 | +// |
| 3 | +// Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +// you may not use this file except in compliance with the License. |
| 5 | +// You may obtain a copy of the License at |
| 6 | +// |
| 7 | +// http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +// |
| 9 | +// Unless required by applicable law or agreed to in writing, software |
| 10 | +// distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +// See the License for the specific language governing permissions and |
| 12 | +// limitations under the License. |
| 13 | + |
| 14 | +package main |
| 15 | + |
| 16 | +import ( |
| 17 | + "fmt" |
| 18 | + "net/http" |
| 19 | + |
| 20 | + "github.com/go-oauth2/oauth2/v4/errors" |
| 21 | + "github.com/go-oauth2/oauth2/v4/generates" |
| 22 | + "github.com/go-oauth2/oauth2/v4/manage" |
| 23 | + "github.com/go-oauth2/oauth2/v4/models" |
| 24 | + "github.com/go-oauth2/oauth2/v4/server" |
| 25 | + "github.com/go-oauth2/oauth2/v4/store" |
| 26 | + "github.com/golang-jwt/jwt" |
| 27 | + "github.com/pingcap/log" |
| 28 | + "github.com/spf13/cobra" |
| 29 | + "go.uber.org/zap" |
| 30 | +) |
| 31 | + |
| 32 | +const openIDConfiguration = ` |
| 33 | +{ |
| 34 | + "issuer": "http://localhost:%d/", |
| 35 | + "authorization_endpoint": "http://localhost:%d/authorize", |
| 36 | + "token_endpoint": "http://localhost:%d/token", |
| 37 | + "scopes_supported": [ |
| 38 | + "openid", |
| 39 | + "profile" |
| 40 | + ], |
| 41 | + "response_types_supported": [ |
| 42 | + "code", |
| 43 | + "token", |
| 44 | + "id_token", |
| 45 | + "code token", |
| 46 | + "code id_token", |
| 47 | + "token id_token", |
| 48 | + "code token id_token" |
| 49 | + ], |
| 50 | + "code_challenge_methods_supported": [ |
| 51 | + "HS256", |
| 52 | + "plain" |
| 53 | + ], |
| 54 | + "response_modes_supported": [ |
| 55 | + "query", |
| 56 | + "fragment", |
| 57 | + "form_post" |
| 58 | + ], |
| 59 | + "subject_types_supported": [ |
| 60 | + "public" |
| 61 | + ], |
| 62 | + "id_token_signing_alg_values_supported": [ |
| 63 | + "HS256" |
| 64 | + ], |
| 65 | + "token_endpoint_auth_methods_supported": [ |
| 66 | + "client_secret_basic", |
| 67 | + "client_secret_post", |
| 68 | + "private_key_jwt" |
| 69 | + ], |
| 70 | + "claims_supported": [ |
| 71 | + "aud", |
| 72 | + "created_at", |
| 73 | + "email", |
| 74 | + "email_verified", |
| 75 | + "exp", |
| 76 | + "iat", |
| 77 | + "identities", |
| 78 | + "iss", |
| 79 | + "name" |
| 80 | + ], |
| 81 | + "request_uri_parameter_supported": false, |
| 82 | + "request_parameter_supported": false |
| 83 | +} |
| 84 | +` |
| 85 | + |
| 86 | +type oauth2ServerConfig struct { |
| 87 | + tokenSignSecret string |
| 88 | + clientID string |
| 89 | + clientSecret string |
| 90 | + port int |
| 91 | +} |
| 92 | + |
| 93 | +var serverConfig = newServerConfig() |
| 94 | + |
| 95 | +func newServerConfig() *oauth2ServerConfig { |
| 96 | + return &oauth2ServerConfig{} |
| 97 | +} |
| 98 | + |
| 99 | +func main() { |
| 100 | + cmd := &cobra.Command{ |
| 101 | + Use: "oauth2 server", |
| 102 | + Run: run, |
| 103 | + } |
| 104 | + // Flags for the root command |
| 105 | + cmd.Flags().StringVar(&serverConfig.tokenSignSecret, "token-secret", |
| 106 | + "SJhX36_KapYSybBtJq35lxX_Brr4LRURSkm7QmXJGmy8pUFW9EIOcVQPsykz9-jj", "Secret to sign token") |
| 107 | + cmd.Flags().StringVar(&serverConfig.clientID, "client-id", "1234", "Client ID of oauth2") |
| 108 | + cmd.Flags().StringVar(&serverConfig.clientSecret, "client-secret", "e0KVlA2EiBfjoN13olyZd2kv1KL", "Client secret of oauth2") |
| 109 | + cmd.Flags().IntVar(&serverConfig.port, "log-file", 9096, "log file path") |
| 110 | + if err := cmd.Execute(); err != nil { |
| 111 | + fmt.Println(err) |
| 112 | + } |
| 113 | +} |
| 114 | + |
| 115 | +func run(_ *cobra.Command, _ []string) { |
| 116 | + manager := manage.NewDefaultManager() |
| 117 | + // token memory store |
| 118 | + manager.MustTokenStorage(store.NewMemoryTokenStore()) |
| 119 | + manager.MapAccessGenerate(generates.NewJWTAccessGenerate("", []byte(serverConfig.tokenSignSecret), jwt.SigningMethodHS512)) |
| 120 | + |
| 121 | + // client memory store |
| 122 | + clientStore := store.NewClientStore() |
| 123 | + err := clientStore.Set(serverConfig.clientID, &models.Client{ |
| 124 | + ID: serverConfig.clientID, |
| 125 | + Secret: serverConfig.clientSecret, |
| 126 | + Domain: fmt.Sprintf("http://localhost:%d", serverConfig.port), |
| 127 | + }) |
| 128 | + if err != nil { |
| 129 | + log.Panic("set client failed", zap.Error(err)) |
| 130 | + } |
| 131 | + manager.MapClientStorage(clientStore) |
| 132 | + |
| 133 | + srv := server.NewDefaultServer(manager) |
| 134 | + srv.SetAllowGetAccessRequest(true) |
| 135 | + srv.SetClientInfoHandler(server.ClientFormHandler) |
| 136 | + srv.SetInternalErrorHandler(func(err error) (re *errors.Response) { |
| 137 | + log.Error("Internal Error:", zap.Error(err)) |
| 138 | + return |
| 139 | + }) |
| 140 | + srv.SetResponseErrorHandler(func(re *errors.Response) { |
| 141 | + log.Error("Response Error:", zap.Error(re.Error)) |
| 142 | + }) |
| 143 | + http.Handle("/authorize", logMiddleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 144 | + err := srv.HandleAuthorizeRequest(w, r) |
| 145 | + if err != nil { |
| 146 | + http.Error(w, err.Error(), http.StatusBadRequest) |
| 147 | + } |
| 148 | + }))) |
| 149 | + http.Handle("/token", logMiddleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 150 | + if err := srv.HandleTokenRequest(w, r); err != nil { |
| 151 | + http.Error(w, err.Error(), http.StatusInternalServerError) |
| 152 | + } |
| 153 | + }))) |
| 154 | + http.Handle("/.well-known/openid-configuration", logMiddleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 155 | + _, _ = w.Write([]byte(fmt.Sprintf(openIDConfiguration, serverConfig.port, serverConfig.port, serverConfig.port))) |
| 156 | + w.WriteHeader(200) |
| 157 | + }))) |
| 158 | + log.Info("starting auth2 server", zap.Int("port", serverConfig.port)) |
| 159 | + log.Panic("run auth2 server failed", zap.Error(http.ListenAndServe(fmt.Sprintf(":%d", serverConfig.port), nil))) |
| 160 | +} |
| 161 | + |
| 162 | +func logMiddleware(next http.Handler) http.Handler { |
| 163 | + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 164 | + next.ServeHTTP(w, r) |
| 165 | + log.Info("oauth server api is called", zap.String("path", r.URL.Path)) |
| 166 | + }) |
| 167 | +} |
0 commit comments