1
0

oauth2.go 6.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249
  1. package sso
  2. import (
  3. "context"
  4. _ "embed"
  5. "encoding/json"
  6. "log"
  7. "net/http"
  8. "net/url"
  9. "time"
  10. "github.com/go-oauth2/oauth2/v4/errors"
  11. "github.com/go-oauth2/oauth2/v4/generates"
  12. "github.com/go-oauth2/oauth2/v4/manage"
  13. "github.com/go-oauth2/oauth2/v4/models"
  14. "github.com/go-oauth2/oauth2/v4/server"
  15. "github.com/go-oauth2/oauth2/v4/store"
  16. "github.com/go-session/session"
  17. "imuslab.com/zoraxy/mod/utils"
  18. )
  19. const (
  20. SSO_SESSION_NAME = "ZoraxySSO"
  21. )
  22. type OAuth2Server struct {
  23. srv *server.Server //oAuth server instance
  24. config *SSOConfig
  25. parent *SSOHandler
  26. }
  27. //go:embed static/auth.html
  28. var authHtml []byte
  29. //go:embed static/login.html
  30. var loginHtml []byte
  31. // NewOAuth2Server creates a new OAuth2 server instance
  32. func NewOAuth2Server(config *SSOConfig, parent *SSOHandler) (*OAuth2Server, error) {
  33. manager := manage.NewDefaultManager()
  34. manager.SetAuthorizeCodeTokenCfg(manage.DefaultAuthorizeCodeTokenCfg)
  35. // token store
  36. manager.MustTokenStorage(store.NewFileTokenStore("./conf/sso.db"))
  37. // generate jwt access token
  38. manager.MapAccessGenerate(generates.NewAccessGenerate())
  39. //Load the information of registered app within the OAuth2 server
  40. clientStore := store.NewClientStore()
  41. clientStore.Set("myapp", &models.Client{
  42. ID: "myapp",
  43. Secret: "verysecurepassword",
  44. Domain: "localhost:9094",
  45. })
  46. //TODO: LOAD THIS DYNAMICALLY FROM DATABASE
  47. manager.MapClientStorage(clientStore)
  48. thisServer := OAuth2Server{
  49. config: config,
  50. parent: parent,
  51. }
  52. //Create a new oauth server
  53. srv := server.NewServer(server.NewConfig(), manager)
  54. srv.SetPasswordAuthorizationHandler(thisServer.PasswordAuthorizationHandler)
  55. srv.SetUserAuthorizationHandler(thisServer.UserAuthorizeHandler)
  56. srv.SetInternalErrorHandler(func(err error) (re *errors.Response) {
  57. log.Println("Internal Error:", err.Error())
  58. return
  59. })
  60. srv.SetResponseErrorHandler(func(re *errors.Response) {
  61. log.Println("Response Error:", re.Error.Error())
  62. })
  63. srv.SetAccessTokenExpHandler(thisServer.ExpireHandler)
  64. thisServer.srv = srv
  65. return &thisServer, nil
  66. }
  67. // Password handler, validate if the given username and password are correct
  68. func (oas *OAuth2Server) PasswordAuthorizationHandler(ctx context.Context, clientID, username, password string) (userID string, err error) {
  69. //TODO: LOAD THIS DYNAMICALLY FROM DATABASE
  70. if username == "test" && password == "test" {
  71. userID = "test"
  72. }
  73. return
  74. }
  75. // User Authorization Handler, handle auth request from user
  76. func (oas *OAuth2Server) UserAuthorizeHandler(w http.ResponseWriter, r *http.Request) (userID string, err error) {
  77. store, err := session.Start(r.Context(), w, r)
  78. if err != nil {
  79. return
  80. }
  81. uid, ok := store.Get(SSO_SESSION_NAME)
  82. if !ok {
  83. if r.Form == nil {
  84. r.ParseForm()
  85. }
  86. store.Set("ReturnUri", r.Form)
  87. store.Save()
  88. w.Header().Set("Location", "/oauth2/login")
  89. w.WriteHeader(http.StatusFound)
  90. return
  91. }
  92. userID = uid.(string)
  93. store.Delete(SSO_SESSION_NAME)
  94. store.Save()
  95. return
  96. }
  97. // AccessTokenExpHandler, set the SSO session length default value
  98. func (oas *OAuth2Server) ExpireHandler(w http.ResponseWriter, r *http.Request) (exp time.Duration, err error) {
  99. requestHostname := r.Host
  100. if requestHostname == "" {
  101. //Use default value
  102. return time.Hour, nil
  103. }
  104. //Get the Registered App Config from parent
  105. appConfig, ok := oas.parent.Apps[requestHostname]
  106. if !ok {
  107. //Use default value
  108. return time.Hour, nil
  109. }
  110. //Use the app's session length
  111. return time.Second * time.Duration(appConfig.SessionDuration), nil
  112. }
  113. /* SSO Web Server Toggle Functions */
  114. func (oas *OAuth2Server) RegisterOauthEndpoints(primaryMux *http.ServeMux) {
  115. primaryMux.HandleFunc("/oauth2/login", oas.loginHandler)
  116. primaryMux.HandleFunc("/oauth2/auth", oas.authHandler)
  117. primaryMux.HandleFunc("/oauth2/authorize", func(w http.ResponseWriter, r *http.Request) {
  118. store, err := session.Start(r.Context(), w, r)
  119. if err != nil {
  120. http.Error(w, err.Error(), http.StatusInternalServerError)
  121. return
  122. }
  123. var form url.Values
  124. if v, ok := store.Get("ReturnUri"); ok {
  125. form = v.(url.Values)
  126. }
  127. r.Form = form
  128. store.Delete("ReturnUri")
  129. store.Save()
  130. err = oas.srv.HandleAuthorizeRequest(w, r)
  131. if err != nil {
  132. http.Error(w, err.Error(), http.StatusBadRequest)
  133. }
  134. })
  135. primaryMux.HandleFunc("/oauth2/token", func(w http.ResponseWriter, r *http.Request) {
  136. err := oas.srv.HandleTokenRequest(w, r)
  137. if err != nil {
  138. http.Error(w, err.Error(), http.StatusInternalServerError)
  139. }
  140. })
  141. primaryMux.HandleFunc("/test", func(w http.ResponseWriter, r *http.Request) {
  142. token, err := oas.srv.ValidationBearerToken(r)
  143. if err != nil {
  144. http.Error(w, err.Error(), http.StatusBadRequest)
  145. return
  146. }
  147. data := map[string]interface{}{
  148. "expires_in": int64(time.Until(token.GetAccessCreateAt().Add(token.GetAccessExpiresIn())).Seconds()),
  149. "client_id": token.GetClientID(),
  150. "user_id": token.GetUserID(),
  151. }
  152. e := json.NewEncoder(w)
  153. e.SetIndent("", " ")
  154. e.Encode(data)
  155. })
  156. }
  157. func (oas *OAuth2Server) loginHandler(w http.ResponseWriter, r *http.Request) {
  158. store, err := session.Start(r.Context(), w, r)
  159. if err != nil {
  160. http.Error(w, err.Error(), http.StatusInternalServerError)
  161. return
  162. }
  163. if r.Method == "POST" {
  164. if r.Form == nil {
  165. if err := r.ParseForm(); err != nil {
  166. http.Error(w, err.Error(), http.StatusInternalServerError)
  167. return
  168. }
  169. }
  170. //Load username and password from form post
  171. username, err := utils.PostPara(r, "username")
  172. if err != nil {
  173. w.Write([]byte("invalid username or password"))
  174. return
  175. }
  176. password, err := utils.PostPara(r, "password")
  177. if err != nil {
  178. w.Write([]byte("invalid username or password"))
  179. return
  180. }
  181. //Validate the user
  182. if !oas.parent.ValidateUsernameAndPassword(username, password) {
  183. //Wrong password
  184. w.Write([]byte("invalid username or password"))
  185. return
  186. }
  187. store.Set(SSO_SESSION_NAME, r.Form.Get("username"))
  188. store.Save()
  189. w.Header().Set("Location", "/oauth2/auth")
  190. w.WriteHeader(http.StatusFound)
  191. return
  192. }
  193. //User not logged in. Show login page
  194. w.Write(loginHtml)
  195. }
  196. func (oas *OAuth2Server) authHandler(w http.ResponseWriter, r *http.Request) {
  197. store, err := session.Start(context.TODO(), w, r)
  198. if err != nil {
  199. http.Error(w, err.Error(), http.StatusInternalServerError)
  200. return
  201. }
  202. if _, ok := store.Get(SSO_SESSION_NAME); !ok {
  203. w.Header().Set("Location", "/oauth2/login")
  204. w.WriteHeader(http.StatusFound)
  205. return
  206. }
  207. //User logged in. Check if this user have previously authorized the app
  208. //TODO: Check if the user have previously authorized the app
  209. //User have not authorized the app. Show the authorization page
  210. w.Write(authHtml)
  211. }