oauth2.go 5.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200
  1. package sso
  2. import (
  3. "context"
  4. _ "embed"
  5. "encoding/json"
  6. "fmt"
  7. "log"
  8. "net/http"
  9. "net/url"
  10. "time"
  11. "github.com/go-oauth2/oauth2/v4/errors"
  12. "github.com/go-oauth2/oauth2/v4/generates"
  13. "github.com/go-oauth2/oauth2/v4/manage"
  14. "github.com/go-oauth2/oauth2/v4/models"
  15. "github.com/go-oauth2/oauth2/v4/server"
  16. "github.com/go-oauth2/oauth2/v4/store"
  17. "github.com/go-session/session"
  18. )
  19. type OAuth2Server struct {
  20. srv *server.Server //oAuth server instance
  21. config *SSOConfig
  22. parent *SSOHandler
  23. }
  24. //go:embed static/auth.html
  25. var authHtml []byte
  26. //go:embed static/login.html
  27. var loginHtml []byte
  28. // NewOAuth2Server creates a new OAuth2 server instance
  29. func NewOAuth2Server(config *SSOConfig, parent *SSOHandler) (*OAuth2Server, error) {
  30. manager := manage.NewDefaultManager()
  31. manager.SetAuthorizeCodeTokenCfg(manage.DefaultAuthorizeCodeTokenCfg)
  32. // token store
  33. manager.MustTokenStorage(store.NewMemoryTokenStore())
  34. // generate jwt access token
  35. manager.MapAccessGenerate(generates.NewAccessGenerate())
  36. //Load user information from SSO_USERS
  37. clientStore := store.NewClientStore()
  38. clientStore.Set("alanyeung", &models.Client{
  39. ID: "alanyeung",
  40. Secret: "password",
  41. Domain: "localhost",
  42. })
  43. manager.MapClientStorage(clientStore)
  44. thisServer := OAuth2Server{
  45. config: config,
  46. parent: parent,
  47. }
  48. //Create a new oauth server
  49. srv := server.NewServer(server.NewConfig(), manager)
  50. srv.SetPasswordAuthorizationHandler(thisServer.PasswordAuthorizationHandler)
  51. srv.SetUserAuthorizationHandler(thisServer.UserAuthorizeHandler)
  52. srv.SetInternalErrorHandler(func(err error) (re *errors.Response) {
  53. log.Println("Internal Error:", err.Error())
  54. return
  55. })
  56. srv.SetResponseErrorHandler(func(re *errors.Response) {
  57. log.Println("Response Error:", re.Error.Error())
  58. })
  59. thisServer.srv = srv
  60. return &thisServer, nil
  61. }
  62. // Password handler, validate if the given username and password are correct
  63. func (oas *OAuth2Server) PasswordAuthorizationHandler(ctx context.Context, clientID, username, password string) (userID string, err error) {
  64. fmt.Println(username, password)
  65. if username == "test" && password == "test" {
  66. userID = "test"
  67. }
  68. return
  69. }
  70. // User Authorization Handler, handle auth request from user
  71. func (oas *OAuth2Server) UserAuthorizeHandler(w http.ResponseWriter, r *http.Request) (userID string, err error) {
  72. store, err := session.Start(r.Context(), w, r)
  73. if err != nil {
  74. return
  75. }
  76. uid, ok := store.Get("ZoraxySSO")
  77. if !ok {
  78. if r.Form == nil {
  79. r.ParseForm()
  80. }
  81. store.Set("ReturnUri", r.Form)
  82. store.Save()
  83. w.Header().Set("Location", "/login")
  84. w.WriteHeader(http.StatusFound)
  85. return
  86. }
  87. userID = uid.(string)
  88. store.Delete("ZoraxySSO")
  89. store.Save()
  90. return
  91. }
  92. /* SSO Web Server Toggle Functions */
  93. func (oas *OAuth2Server) RegisterOauthEndpoints(primaryMux *http.ServeMux) {
  94. primaryMux.HandleFunc("/oauth2/login", loginHandler)
  95. primaryMux.HandleFunc("/oauth2/auth", authHandler)
  96. primaryMux.HandleFunc("/oauth2/authorize", func(w http.ResponseWriter, r *http.Request) {
  97. store, err := session.Start(r.Context(), w, r)
  98. if err != nil {
  99. http.Error(w, err.Error(), http.StatusInternalServerError)
  100. return
  101. }
  102. var form url.Values
  103. if v, ok := store.Get("ReturnUri"); ok {
  104. form = v.(url.Values)
  105. }
  106. r.Form = form
  107. store.Delete("ReturnUri")
  108. store.Save()
  109. err = oas.srv.HandleAuthorizeRequest(w, r)
  110. if err != nil {
  111. http.Error(w, err.Error(), http.StatusBadRequest)
  112. }
  113. })
  114. primaryMux.HandleFunc("/oauth2/token", func(w http.ResponseWriter, r *http.Request) {
  115. err := oas.srv.HandleTokenRequest(w, r)
  116. if err != nil {
  117. http.Error(w, err.Error(), http.StatusInternalServerError)
  118. }
  119. })
  120. primaryMux.HandleFunc("/test", func(w http.ResponseWriter, r *http.Request) {
  121. token, err := oas.srv.ValidationBearerToken(r)
  122. if err != nil {
  123. http.Error(w, err.Error(), http.StatusBadRequest)
  124. return
  125. }
  126. data := map[string]interface{}{
  127. "expires_in": int64(time.Until(token.GetAccessCreateAt().Add(token.GetAccessExpiresIn())).Seconds()),
  128. "client_id": token.GetClientID(),
  129. "user_id": token.GetUserID(),
  130. }
  131. e := json.NewEncoder(w)
  132. e.SetIndent("", " ")
  133. e.Encode(data)
  134. })
  135. }
  136. func loginHandler(w http.ResponseWriter, r *http.Request) {
  137. store, err := session.Start(r.Context(), w, r)
  138. if err != nil {
  139. http.Error(w, err.Error(), http.StatusInternalServerError)
  140. return
  141. }
  142. if r.Method == "POST" {
  143. if r.Form == nil {
  144. if err := r.ParseForm(); err != nil {
  145. http.Error(w, err.Error(), http.StatusInternalServerError)
  146. return
  147. }
  148. }
  149. store.Set("ZoraxySSO", r.Form.Get("username"))
  150. store.Save()
  151. w.Header().Set("Location", "/oauth2/auth")
  152. w.WriteHeader(http.StatusFound)
  153. return
  154. }
  155. //User not logged in. Show login page
  156. w.Write(loginHtml)
  157. }
  158. func authHandler(w http.ResponseWriter, r *http.Request) {
  159. store, err := session.Start(context.TODO(), w, r)
  160. if err != nil {
  161. http.Error(w, err.Error(), http.StatusInternalServerError)
  162. return
  163. }
  164. if _, ok := store.Get("ZoraxySSO"); !ok {
  165. w.Header().Set("Location", "/oauth2/login")
  166. w.WriteHeader(http.StatusFound)
  167. return
  168. }
  169. //User logged in. Show authentication confirm page
  170. w.Write(authHtml)
  171. }