builder.go 36 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422
  1. package login
  2. import (
  3. "context"
  4. "errors"
  5. "fmt"
  6. "io/fs"
  7. "net/http"
  8. "net/url"
  9. "reflect"
  10. "strings"
  11. "time"
  12. "github.com/golang-jwt/jwt/v4"
  13. "github.com/markbates/goth"
  14. "github.com/markbates/goth/gothic"
  15. "github.com/pquerna/otp"
  16. "github.com/pquerna/otp/totp"
  17. "github.com/qor5/web"
  18. "github.com/qor5/x/i18n"
  19. h "github.com/theplant/htmlgo"
  20. "golang.org/x/text/language"
  21. "gorm.io/gorm"
  22. )
  23. var (
  24. ErrUserNotFound = errors.New("user not found")
  25. ErrPasswordChanged = errors.New("password changed")
  26. ErrWrongPassword = errors.New("wrong password")
  27. ErrUserLocked = errors.New("user locked")
  28. ErrUserGetLocked = errors.New("user get locked")
  29. ErrWrongTOTPCode = errors.New("wrong totp code")
  30. ErrTOTPCodeHasBeenUsed = errors.New("totp code has been used")
  31. ErrEmptyPassword = errors.New("empty password")
  32. ErrPasswordNotMatch = errors.New("password not match")
  33. )
  34. type HomeURLFunc func(r *http.Request, user interface{}) string
  35. type HookFunc func(r *http.Request, user interface{}, extraVals ...interface{}) error
  36. type Provider struct {
  37. Goth goth.Provider
  38. Key string
  39. Text string
  40. Logo h.HTMLComponent
  41. }
  42. type CookieConfig struct {
  43. Path string
  44. Domain string
  45. SameSite http.SameSite
  46. }
  47. type TOTPConfig struct {
  48. Issuer string
  49. }
  50. type RecaptchaConfig struct {
  51. SiteKey string
  52. SecretKey string
  53. }
  54. type Builder struct {
  55. secret string
  56. providers []*Provider
  57. authCookieName string
  58. authSecureCookieName string
  59. continueUrlCookieName string
  60. // seconds
  61. sessionMaxAge int
  62. cookieConfig CookieConfig
  63. totpEnabled bool
  64. totpConfig TOTPConfig
  65. recaptchaEnabled bool
  66. recaptchaConfig RecaptchaConfig
  67. autoExtendSession bool
  68. maxRetryCount int
  69. noForgetPasswordLink bool
  70. i18nBuilder *i18n.Builder
  71. // Common URLs
  72. homePageURLFunc HomeURLFunc
  73. loginPageURL string
  74. LogoutURL string
  75. // TOTP URLs
  76. validateTOTPURL string
  77. totpSetupPageURL string
  78. totpValidatePageURL string
  79. // OAuth URLs
  80. oauthBeginURL string
  81. oauthCallbackURL string
  82. oauthCallbackCompleteURL string
  83. // UserPass URLs
  84. passwordLoginURL string
  85. resetPasswordURL string
  86. resetPasswordPageURL string
  87. changePasswordURL string
  88. changePasswordPageURL string
  89. forgetPasswordPageURL string
  90. sendResetPasswordLinkURL string
  91. resetPasswordLinkSentPageURL string
  92. loginPageFunc web.PageFunc
  93. forgetPasswordPageFunc web.PageFunc
  94. resetPasswordLinkSentPageFunc web.PageFunc
  95. resetPasswordPageFunc web.PageFunc
  96. changePasswordPageFunc web.PageFunc
  97. totpSetupPageFunc web.PageFunc
  98. totpValidatePageFunc web.PageFunc
  99. beforeSetPasswordHook HookFunc
  100. afterLoginHook HookFunc
  101. afterFailedToLoginHook HookFunc
  102. afterUserLockedHook HookFunc
  103. afterLogoutHook HookFunc
  104. afterConfirmSendResetPasswordLinkHook HookFunc
  105. afterResetPasswordHook HookFunc
  106. afterChangePasswordHook HookFunc
  107. afterExtendSessionHook HookFunc
  108. afterTOTPCodeReusedHook HookFunc
  109. afterOAuthCompleteHook HookFunc
  110. db *gorm.DB
  111. userModel interface{}
  112. snakePrimaryField string
  113. tUser reflect.Type
  114. userPassEnabled bool
  115. oauthEnabled bool
  116. sessionSecureEnabled bool
  117. }
  118. func New() *Builder {
  119. r := &Builder{
  120. authCookieName: "auth",
  121. authSecureCookieName: "qor5_auth_secure",
  122. continueUrlCookieName: "qor5_continue_url",
  123. homePageURLFunc: func(r *http.Request, user interface{}) string {
  124. return "/"
  125. },
  126. loginPageURL: "/auth/login",
  127. LogoutURL: "/auth/logout",
  128. validateTOTPURL: "/auth/2fa/totp/do",
  129. totpSetupPageURL: "/auth/2fa/totp/setup",
  130. totpValidatePageURL: "/auth/2fa/totp/validate",
  131. oauthBeginURL: "/auth/begin",
  132. oauthCallbackURL: "/auth/callback",
  133. oauthCallbackCompleteURL: "/auth/callback-complete",
  134. passwordLoginURL: "/auth/userpass/login",
  135. resetPasswordURL: "/auth/do-reset-password",
  136. resetPasswordPageURL: "/auth/reset-password",
  137. changePasswordURL: "/auth/do-change-password",
  138. changePasswordPageURL: "/auth/change-password",
  139. forgetPasswordPageURL: "/auth/forget-password",
  140. sendResetPasswordLinkURL: "/auth/send-reset-password-link",
  141. resetPasswordLinkSentPageURL: "/auth/reset-password-link-sent",
  142. sessionMaxAge: 60 * 60,
  143. cookieConfig: CookieConfig{
  144. Path: "/",
  145. Domain: "",
  146. SameSite: http.SameSiteStrictMode,
  147. },
  148. autoExtendSession: true,
  149. maxRetryCount: 5,
  150. totpEnabled: true,
  151. totpConfig: TOTPConfig{
  152. Issuer: "QOR5",
  153. },
  154. }
  155. i18nB := i18n.New()
  156. i18nB.SupportLanguages(language.English, language.SimplifiedChinese, language.Japanese)
  157. r.I18n(i18nB)
  158. vh := r.ViewHelper()
  159. r.loginPageFunc = defaultLoginPage(vh)
  160. r.forgetPasswordPageFunc = defaultForgetPasswordPage(vh)
  161. r.resetPasswordLinkSentPageFunc = defaultResetPasswordLinkSentPage(vh)
  162. r.resetPasswordPageFunc = defaultResetPasswordPage(vh)
  163. r.changePasswordPageFunc = defaultChangePasswordPage(vh)
  164. r.totpSetupPageFunc = defaultTOTPSetupPage(vh)
  165. r.totpValidatePageFunc = defaultTOTPValidatePage(vh)
  166. return r
  167. }
  168. func (b *Builder) Secret(v string) (r *Builder) {
  169. b.secret = v
  170. return b
  171. }
  172. func (b *Builder) CookieConfig(v CookieConfig) (r *Builder) {
  173. b.cookieConfig = v
  174. return b
  175. }
  176. // Google reCAPTCHA.
  177. func (b *Builder) Recaptcha(enable bool, config ...RecaptchaConfig) (r *Builder) {
  178. b.recaptchaEnabled = enable
  179. if len(config) > 0 {
  180. b.recaptchaConfig = config[0]
  181. }
  182. if enable {
  183. if b.recaptchaConfig.SiteKey == "" {
  184. panic("SiteKey is empty")
  185. }
  186. if b.recaptchaConfig.SecretKey == "" {
  187. panic("SecretKey is empty")
  188. }
  189. }
  190. return b
  191. }
  192. func (b *Builder) OAuthProviders(vs ...*Provider) (r *Builder) {
  193. if len(vs) == 0 {
  194. return b
  195. }
  196. b.oauthEnabled = true
  197. b.providers = vs
  198. var gothProviders []goth.Provider
  199. for _, v := range vs {
  200. gothProviders = append(gothProviders, v.Goth)
  201. }
  202. goth.UseProviders(gothProviders...)
  203. return b
  204. }
  205. func (b *Builder) AuthCookieName(v string) (r *Builder) {
  206. b.authCookieName = v
  207. return b
  208. }
  209. func (b *Builder) HomeURLFunc(v HomeURLFunc) (r *Builder) {
  210. b.homePageURLFunc = v
  211. return b
  212. }
  213. func (b *Builder) LoginPageURL(v string) (r *Builder) {
  214. b.loginPageURL = v
  215. return b
  216. }
  217. func (b *Builder) ResetPasswordPageURL(v string) (r *Builder) {
  218. b.resetPasswordPageURL = v
  219. return b
  220. }
  221. func (b *Builder) ChangePasswordPageURL(v string) (r *Builder) {
  222. b.changePasswordPageURL = v
  223. return b
  224. }
  225. func (b *Builder) ForgetPasswordPageURL(v string) (r *Builder) {
  226. b.forgetPasswordPageURL = v
  227. return b
  228. }
  229. func (b *Builder) ResetPasswordLinkSentPageURL(v string) (r *Builder) {
  230. b.resetPasswordLinkSentPageURL = v
  231. return b
  232. }
  233. func (b *Builder) TOTPSetupPageURL(v string) (r *Builder) {
  234. b.totpSetupPageURL = v
  235. return b
  236. }
  237. func (b *Builder) TOTPValidatePageURL(v string) (r *Builder) {
  238. b.totpValidatePageURL = v
  239. return b
  240. }
  241. func (b *Builder) LoginPageFunc(v web.PageFunc) (r *Builder) {
  242. b.loginPageFunc = v
  243. return b
  244. }
  245. func (b *Builder) ForgetPasswordPageFunc(v web.PageFunc) (r *Builder) {
  246. b.forgetPasswordPageFunc = v
  247. return b
  248. }
  249. func (b *Builder) ResetPasswordLinkSentPageFunc(v web.PageFunc) (r *Builder) {
  250. b.resetPasswordLinkSentPageFunc = v
  251. return b
  252. }
  253. func (b *Builder) ResetPasswordPageFunc(v web.PageFunc) (r *Builder) {
  254. b.resetPasswordPageFunc = v
  255. return b
  256. }
  257. func (b *Builder) ChangePasswordPageFunc(v web.PageFunc) (r *Builder) {
  258. b.changePasswordPageFunc = v
  259. return b
  260. }
  261. func (b *Builder) TOTPSetupPageFunc(v web.PageFunc) (r *Builder) {
  262. b.totpSetupPageFunc = v
  263. return b
  264. }
  265. func (b *Builder) TOTPValidatePageFunc(v web.PageFunc) (r *Builder) {
  266. b.totpValidatePageFunc = v
  267. return b
  268. }
  269. func (b *Builder) wrapHook(v HookFunc) HookFunc {
  270. if v == nil {
  271. return nil
  272. }
  273. return func(r *http.Request, user interface{}, extraVals ...interface{}) error {
  274. if user != nil && GetCurrentUser(r) == nil {
  275. r = r.WithContext(context.WithValue(r.Context(), UserKey, user))
  276. }
  277. return v(r, user, extraVals...)
  278. }
  279. }
  280. // extra vals:
  281. // - password
  282. func (b *Builder) BeforeSetPassword(v HookFunc) (r *Builder) {
  283. b.beforeSetPasswordHook = b.wrapHook(v)
  284. return b
  285. }
  286. func (b *Builder) AfterLogin(v HookFunc) (r *Builder) {
  287. b.afterLoginHook = b.wrapHook(v)
  288. return b
  289. }
  290. // extra vals:
  291. // - login error
  292. func (b *Builder) AfterFailedToLogin(v HookFunc) (r *Builder) {
  293. b.afterFailedToLoginHook = b.wrapHook(v)
  294. return b
  295. }
  296. func (b *Builder) AfterUserLocked(v HookFunc) (r *Builder) {
  297. b.afterUserLockedHook = b.wrapHook(v)
  298. return b
  299. }
  300. func (b *Builder) AfterLogout(v HookFunc) (r *Builder) {
  301. b.afterLogoutHook = b.wrapHook(v)
  302. return b
  303. }
  304. // extra vals:
  305. // - reset link
  306. func (b *Builder) AfterConfirmSendResetPasswordLink(v HookFunc) (r *Builder) {
  307. b.afterConfirmSendResetPasswordLinkHook = b.wrapHook(v)
  308. return b
  309. }
  310. func (b *Builder) AfterResetPassword(v HookFunc) (r *Builder) {
  311. b.afterResetPasswordHook = b.wrapHook(v)
  312. return b
  313. }
  314. func (b *Builder) AfterChangePassword(v HookFunc) (r *Builder) {
  315. b.afterChangePasswordHook = b.wrapHook(v)
  316. return b
  317. }
  318. // extra vals:
  319. // - old session token
  320. func (b *Builder) AfterExtendSession(v HookFunc) (r *Builder) {
  321. b.afterExtendSessionHook = b.wrapHook(v)
  322. return b
  323. }
  324. func (b *Builder) AfterTOTPCodeReused(v HookFunc) (r *Builder) {
  325. b.afterTOTPCodeReusedHook = b.wrapHook(v)
  326. return b
  327. }
  328. // user is goth.User
  329. func (b *Builder) AfterOAuthComplete(v HookFunc) (r *Builder) {
  330. b.afterOAuthCompleteHook = b.wrapHook(v)
  331. return b
  332. }
  333. // seconds
  334. // default 1h
  335. func (b *Builder) SessionMaxAge(v int) (r *Builder) {
  336. b.sessionMaxAge = v
  337. return b
  338. }
  339. // extend the session if successfully authenticated
  340. // default true
  341. func (b *Builder) AutoExtendSession(v bool) (r *Builder) {
  342. b.autoExtendSession = v
  343. return b
  344. }
  345. // default 5
  346. // MaxRetryCount <= 0 means no max retry count limit
  347. func (b *Builder) MaxRetryCount(v int) (r *Builder) {
  348. b.maxRetryCount = v
  349. return b
  350. }
  351. func (b *Builder) TOTP(enable bool, config ...TOTPConfig) (r *Builder) {
  352. b.totpEnabled = enable
  353. if len(config) > 0 {
  354. b.totpConfig = config[0]
  355. }
  356. if enable {
  357. if b.totpConfig.Issuer == "" {
  358. panic("Issuer is empty")
  359. }
  360. }
  361. return b
  362. }
  363. func (b *Builder) NoForgetPasswordLink(v bool) (r *Builder) {
  364. b.noForgetPasswordLink = v
  365. return b
  366. }
  367. func (b *Builder) DB(v *gorm.DB) (r *Builder) {
  368. b.db = v
  369. return b
  370. }
  371. func (b *Builder) I18n(v *i18n.Builder) (r *Builder) {
  372. v.RegisterForModule(language.English, I18nLoginKey, Messages_en_US).
  373. RegisterForModule(language.SimplifiedChinese, I18nLoginKey, Messages_zh_CN).
  374. RegisterForModule(language.Japanese, I18nLoginKey, Messages_ja_JP)
  375. b.i18nBuilder = v
  376. return b
  377. }
  378. func (b *Builder) GetSessionMaxAge() int {
  379. return b.sessionMaxAge
  380. }
  381. func (b *Builder) ViewHelper() *ViewHelper {
  382. return &ViewHelper{
  383. b: b,
  384. }
  385. }
  386. func (b *Builder) UserModel(m interface{}) (r *Builder) {
  387. b.userModel = m
  388. b.tUser = underlyingReflectType(reflect.TypeOf(m))
  389. b.snakePrimaryField = snakePrimaryField(m)
  390. if _, ok := m.(UserPasser); ok {
  391. b.userPassEnabled = true
  392. }
  393. if _, ok := m.(OAuthUser); ok {
  394. b.oauthEnabled = true
  395. }
  396. if _, ok := m.(SessionSecurer); ok {
  397. b.sessionSecureEnabled = true
  398. }
  399. return b
  400. }
  401. func (b *Builder) newUserObject() interface{} {
  402. return reflect.New(b.tUser).Interface()
  403. }
  404. func (b *Builder) findUserByID(id string) (user interface{}, err error) {
  405. m := b.newUserObject()
  406. err = b.db.Where(fmt.Sprintf("%s = ?", b.snakePrimaryField), id).
  407. First(m).
  408. Error
  409. if err != nil {
  410. if err == gorm.ErrRecordNotFound {
  411. return nil, ErrUserNotFound
  412. }
  413. return nil, err
  414. }
  415. return m, nil
  416. }
  417. // completeUserAuthCallback is for url "/auth/{provider}/callback"
  418. func (b *Builder) completeUserAuthCallback(w http.ResponseWriter, r *http.Request) {
  419. if b.cookieConfig.SameSite != http.SameSiteStrictMode {
  420. b.completeUserAuthCallbackComplete(w, r)
  421. return
  422. }
  423. completeURL := fmt.Sprintf("%s?%s", b.oauthCallbackCompleteURL, r.URL.Query().Encode())
  424. w.Header().Set("Content-Type", "text/html; charset=utf-8")
  425. w.Write([]byte(fmt.Sprintf(`
  426. <script>
  427. window.location.href="%s";
  428. </script>
  429. <a href="%s">complete</a>
  430. `, completeURL, completeURL)))
  431. return
  432. }
  433. func (b *Builder) completeUserAuthCallbackComplete(w http.ResponseWriter, r *http.Request) {
  434. var err error
  435. var user interface{}
  436. failRedirectURL := b.LogoutURL
  437. defer func() {
  438. if err != nil {
  439. if b.afterFailedToLoginHook != nil {
  440. if herr := b.afterFailedToLoginHook(r, user, err); herr != nil {
  441. setNoticeOrFailCodeFlash(w, herr, FailCodeSystemError)
  442. }
  443. }
  444. http.Redirect(w, r, failRedirectURL, http.StatusFound)
  445. }
  446. }()
  447. var ouser goth.User
  448. ouser, err = gothic.CompleteUserAuth(w, r)
  449. if err != nil {
  450. setFailCodeFlash(w, FailCodeCompleteUserAuthFailed)
  451. return
  452. }
  453. if b.afterOAuthCompleteHook != nil {
  454. if err = b.afterOAuthCompleteHook(r, ouser); err != nil {
  455. setNoticeOrFailCodeFlash(w, err, FailCodeSystemError)
  456. return
  457. }
  458. }
  459. userID := ouser.UserID
  460. if b.userModel != nil {
  461. user, err = b.userModel.(OAuthUser).FindUserByOAuthUserID(b.db, b.newUserObject(), ouser.Provider, ouser.UserID)
  462. if err != nil {
  463. if err != gorm.ErrRecordNotFound {
  464. setFailCodeFlash(w, FailCodeSystemError)
  465. return
  466. }
  467. // TODO: maybe the identifier of some providers is not email
  468. identifier := ouser.Email
  469. user, err = b.userModel.(OAuthUser).FindUserByOAuthIdentifier(b.db, b.newUserObject(), ouser.Provider, identifier)
  470. if err != nil {
  471. if err == gorm.ErrRecordNotFound {
  472. setFailCodeFlash(w, FailCodeUserNotFound)
  473. } else {
  474. setFailCodeFlash(w, FailCodeSystemError)
  475. }
  476. return
  477. }
  478. err = user.(OAuthUser).InitOAuthUserID(b.db, b.newUserObject(), ouser.Provider, identifier, ouser.UserID)
  479. if err != nil {
  480. setFailCodeFlash(w, FailCodeSystemError)
  481. return
  482. }
  483. }
  484. userID = objectID(user)
  485. }
  486. claims := UserClaims{
  487. Provider: ouser.Provider,
  488. Email: ouser.Email,
  489. Name: ouser.Name,
  490. UserID: userID,
  491. AvatarURL: ouser.AvatarURL,
  492. RegisteredClaims: b.genBaseSessionClaim(userID),
  493. }
  494. if user == nil {
  495. user = &claims
  496. }
  497. if b.afterLoginHook != nil {
  498. setCookieForRequest(r, &http.Cookie{Name: b.authCookieName, Value: b.mustGetSessionToken(claims)})
  499. if err = b.afterLoginHook(r, user); err != nil {
  500. setNoticeOrFailCodeFlash(w, err, FailCodeSystemError)
  501. return
  502. }
  503. }
  504. if err = b.setSecureCookiesByClaims(w, user, claims); err != nil {
  505. setFailCodeFlash(w, FailCodeSystemError)
  506. return
  507. }
  508. redirectURL := b.homePageURLFunc(r, user)
  509. if v := b.getContinueURL(w, r); v != "" {
  510. redirectURL = v
  511. }
  512. http.Redirect(w, r, redirectURL, http.StatusFound)
  513. return
  514. }
  515. // return user if account exists even if there is an error returned
  516. func (b *Builder) authUserPass(account string, password string) (user interface{}, err error) {
  517. user, err = b.userModel.(UserPasser).FindUser(b.db, b.newUserObject(), account)
  518. if err != nil {
  519. if err == gorm.ErrRecordNotFound {
  520. return nil, ErrUserNotFound
  521. }
  522. return nil, err
  523. }
  524. u := user.(UserPasser)
  525. if u.GetLocked() {
  526. return user, ErrUserLocked
  527. }
  528. if !u.IsPasswordCorrect(password) {
  529. if b.maxRetryCount > 0 {
  530. if err = u.IncreaseRetryCount(b.db, b.newUserObject()); err != nil {
  531. return user, err
  532. }
  533. if u.GetLoginRetryCount() >= b.maxRetryCount {
  534. if err = u.LockUser(b.db, b.newUserObject()); err != nil {
  535. return user, err
  536. }
  537. return user, ErrUserGetLocked
  538. }
  539. }
  540. return user, ErrWrongPassword
  541. }
  542. if u.GetLoginRetryCount() != 0 {
  543. if err = u.UnlockUser(b.db, b.newUserObject()); err != nil {
  544. return user, err
  545. }
  546. }
  547. return user, nil
  548. }
  549. func (b *Builder) userpassLogin(w http.ResponseWriter, r *http.Request) {
  550. if r.Method != http.MethodPost {
  551. w.WriteHeader(http.StatusNotFound)
  552. return
  553. }
  554. // check reCAPTCHA token
  555. if b.recaptchaEnabled {
  556. token := r.FormValue("token")
  557. if !recaptchaTokenCheck(b, token) {
  558. setFailCodeFlash(w, FailCodeIncorrectRecaptchaToken)
  559. http.Redirect(w, r, b.loginPageURL, http.StatusFound)
  560. return
  561. }
  562. }
  563. var err error
  564. var user interface{}
  565. failRedirectURL := b.LogoutURL
  566. defer func() {
  567. if err != nil {
  568. if b.afterFailedToLoginHook != nil {
  569. if herr := b.afterFailedToLoginHook(r, user, err); herr != nil {
  570. setNoticeOrFailCodeFlash(w, herr, FailCodeSystemError)
  571. }
  572. }
  573. http.Redirect(w, r, failRedirectURL, http.StatusFound)
  574. }
  575. }()
  576. account := r.FormValue("account")
  577. password := r.FormValue("password")
  578. user, err = b.authUserPass(account, password)
  579. if err != nil {
  580. if err == ErrUserGetLocked && b.afterUserLockedHook != nil {
  581. if err = b.afterUserLockedHook(r, user); err != nil {
  582. setNoticeOrFailCodeFlash(w, err, FailCodeSystemError)
  583. return
  584. }
  585. }
  586. code := FailCodeSystemError
  587. switch err {
  588. case ErrWrongPassword, ErrUserNotFound:
  589. code = FailCodeIncorrectAccountNameOrPassword
  590. case ErrUserLocked, ErrUserGetLocked:
  591. code = FailCodeUserLocked
  592. }
  593. setFailCodeFlash(w, code)
  594. setWrongLoginInputFlash(w, WrongLoginInputFlash{
  595. Account: account,
  596. Password: password,
  597. })
  598. return
  599. }
  600. u := user.(UserPasser)
  601. userID := objectID(user)
  602. claims := UserClaims{
  603. UserID: userID,
  604. PassUpdatedAt: u.GetPasswordUpdatedAt(),
  605. RegisteredClaims: b.genBaseSessionClaim(userID),
  606. }
  607. if !b.totpEnabled {
  608. if b.afterLoginHook != nil {
  609. setCookieForRequest(r, &http.Cookie{Name: b.authCookieName, Value: b.mustGetSessionToken(claims)})
  610. if err = b.afterLoginHook(r, user); err != nil {
  611. setNoticeOrFailCodeFlash(w, err, FailCodeSystemError)
  612. return
  613. }
  614. }
  615. }
  616. if err = b.setSecureCookiesByClaims(w, user, claims); err != nil {
  617. setFailCodeFlash(w, FailCodeSystemError)
  618. return
  619. }
  620. if b.totpEnabled {
  621. if u.GetIsTOTPSetup() {
  622. http.Redirect(w, r, b.totpValidatePageURL, http.StatusFound)
  623. return
  624. }
  625. var key *otp.Key
  626. if key, err = totp.Generate(
  627. totp.GenerateOpts{
  628. Issuer: b.totpConfig.Issuer,
  629. AccountName: u.GetAccountName(),
  630. },
  631. ); err != nil {
  632. setFailCodeFlash(w, FailCodeSystemError)
  633. return
  634. }
  635. if err = u.SetTOTPSecret(b.db, b.newUserObject(), key.Secret()); err != nil {
  636. setFailCodeFlash(w, FailCodeSystemError)
  637. return
  638. }
  639. http.Redirect(w, r, b.totpSetupPageURL, http.StatusFound)
  640. return
  641. }
  642. redirectURL := b.homePageURLFunc(r, user)
  643. if v := b.getContinueURL(w, r); v != "" {
  644. redirectURL = v
  645. }
  646. http.Redirect(w, r, redirectURL, http.StatusFound)
  647. return
  648. }
  649. func (b *Builder) genBaseSessionClaim(id string) jwt.RegisteredClaims {
  650. return genBaseClaims(id, b.sessionMaxAge)
  651. }
  652. func (b *Builder) mustGetSessionToken(claims UserClaims) string {
  653. return mustSignClaims(claims, b.secret)
  654. }
  655. func (b *Builder) setAuthCookiesFromUserClaims(w http.ResponseWriter, claims *UserClaims, secureSalt string) error {
  656. http.SetCookie(w, &http.Cookie{
  657. Name: b.authCookieName,
  658. Value: b.mustGetSessionToken(*claims),
  659. Path: b.cookieConfig.Path,
  660. Domain: b.cookieConfig.Domain,
  661. MaxAge: b.sessionMaxAge,
  662. Expires: time.Now().Add(time.Duration(b.sessionMaxAge) * time.Second),
  663. HttpOnly: true,
  664. Secure: true,
  665. SameSite: b.cookieConfig.SameSite,
  666. })
  667. if secureSalt != "" {
  668. http.SetCookie(w, &http.Cookie{
  669. Name: b.authSecureCookieName,
  670. Value: mustSignClaims(&claims.RegisteredClaims, b.secret+secureSalt),
  671. Path: b.cookieConfig.Path,
  672. Domain: b.cookieConfig.Domain,
  673. MaxAge: b.sessionMaxAge,
  674. Expires: time.Now().Add(time.Duration(b.sessionMaxAge) * time.Second),
  675. HttpOnly: true,
  676. Secure: true,
  677. SameSite: b.cookieConfig.SameSite,
  678. })
  679. }
  680. return nil
  681. }
  682. func (b *Builder) cleanAuthCookies(w http.ResponseWriter) {
  683. http.SetCookie(w, &http.Cookie{
  684. Name: b.authCookieName,
  685. Value: "",
  686. Path: b.cookieConfig.Path,
  687. Domain: b.cookieConfig.Domain,
  688. MaxAge: -1,
  689. Expires: time.Unix(1, 0),
  690. HttpOnly: true,
  691. Secure: true,
  692. })
  693. http.SetCookie(w, &http.Cookie{
  694. Name: b.authSecureCookieName,
  695. Value: "",
  696. Path: b.cookieConfig.Path,
  697. Domain: b.cookieConfig.Domain,
  698. MaxAge: -1,
  699. Expires: time.Unix(1, 0),
  700. HttpOnly: true,
  701. Secure: true,
  702. })
  703. }
  704. func (b *Builder) setContinueURL(w http.ResponseWriter, r *http.Request) {
  705. continueURL := r.RequestURI
  706. if strings.Contains(continueURL, "?__execute_event__=") {
  707. continueURL = r.Referer()
  708. }
  709. ignore := false
  710. {
  711. ignoreURLs := map[string]struct{}{
  712. b.loginPageURL: {},
  713. b.resetPasswordPageURL: {},
  714. b.forgetPasswordPageURL: {},
  715. b.resetPasswordLinkSentPageURL: {},
  716. b.totpSetupPageURL: {},
  717. b.totpValidatePageURL: {},
  718. b.LogoutURL: {},
  719. }
  720. u, err := url.Parse(continueURL)
  721. if err != nil {
  722. ignore = true
  723. } else {
  724. if _, ok := ignoreURLs[u.Path]; ok {
  725. ignore = true
  726. }
  727. }
  728. }
  729. if ignore {
  730. return
  731. }
  732. http.SetCookie(w, &http.Cookie{
  733. Name: b.continueUrlCookieName,
  734. Value: continueURL,
  735. Path: b.cookieConfig.Path,
  736. Domain: b.cookieConfig.Domain,
  737. HttpOnly: true,
  738. })
  739. }
  740. func (b *Builder) getContinueURL(w http.ResponseWriter, r *http.Request) string {
  741. c, err := r.Cookie(b.continueUrlCookieName)
  742. if err != nil || c.Value == "" {
  743. return ""
  744. }
  745. http.SetCookie(w, &http.Cookie{
  746. Name: b.continueUrlCookieName,
  747. Value: "",
  748. MaxAge: -1,
  749. Expires: time.Unix(1, 0),
  750. Path: b.cookieConfig.Path,
  751. Domain: b.cookieConfig.Domain,
  752. HttpOnly: true,
  753. })
  754. return c.Value
  755. }
  756. func (b *Builder) setSecureCookiesByClaims(w http.ResponseWriter, user interface{}, claims UserClaims) (err error) {
  757. var secureSalt string
  758. if b.sessionSecureEnabled {
  759. if user.(SessionSecurer).GetSecure() == "" {
  760. err = user.(SessionSecurer).UpdateSecure(b.db, b.newUserObject(), objectID(user))
  761. if err != nil {
  762. return err
  763. }
  764. }
  765. secureSalt = user.(SessionSecurer).GetSecure()
  766. }
  767. if err = b.setAuthCookiesFromUserClaims(w, &claims, secureSalt); err != nil {
  768. return err
  769. }
  770. return nil
  771. }
  772. func (b *Builder) consumeTOTPCode(r *http.Request, up UserPasser, passcode string) error {
  773. if !totp.Validate(passcode, up.GetTOTPSecret()) {
  774. return ErrWrongTOTPCode
  775. }
  776. lastCode, usedAt := up.GetLastUsedTOTPCode()
  777. if usedAt != nil && time.Now().Sub(*usedAt) > 90*time.Second {
  778. lastCode = ""
  779. }
  780. if passcode == lastCode {
  781. if b.afterTOTPCodeReusedHook != nil {
  782. if herr := b.afterTOTPCodeReusedHook(r, GetCurrentUser(r)); herr != nil {
  783. return herr
  784. }
  785. }
  786. return ErrTOTPCodeHasBeenUsed
  787. }
  788. if err := up.SetLastUsedTOTPCode(b.db, b.newUserObject(), passcode); err != nil {
  789. return err
  790. }
  791. return nil
  792. }
  793. func (b *Builder) getFailCodeFromTOTPCodeConsumeError(verr error) FailCode {
  794. fc := FailCodeSystemError
  795. switch verr {
  796. case ErrWrongTOTPCode:
  797. fc = FailCodeIncorrectTOTPCode
  798. case ErrTOTPCodeHasBeenUsed:
  799. fc = FailCodeTOTPCodeHasBeenUsed
  800. }
  801. return fc
  802. }
  803. // logout is for url "/logout/{provider}"
  804. func (b *Builder) logout(w http.ResponseWriter, r *http.Request) {
  805. err := gothic.Logout(w, r)
  806. if err != nil {
  807. //
  808. }
  809. b.cleanAuthCookies(w)
  810. if b.afterLogoutHook != nil {
  811. user := GetCurrentUser(r)
  812. if user != nil {
  813. if herr := b.afterLogoutHook(r, user); herr != nil {
  814. setNoticeOrFailCodeFlash(w, herr, FailCodeSystemError)
  815. http.Redirect(w, r, b.loginPageURL, http.StatusFound)
  816. return
  817. }
  818. }
  819. }
  820. http.Redirect(w, r, b.loginPageURL, http.StatusFound)
  821. }
  822. // beginAuth is for url "/auth/{provider}"
  823. func (b *Builder) beginAuth(w http.ResponseWriter, r *http.Request) {
  824. gothic.BeginAuthHandler(w, r)
  825. }
  826. func (b *Builder) sendResetPasswordLink(w http.ResponseWriter, r *http.Request) {
  827. if r.Method != http.MethodPost {
  828. w.WriteHeader(http.StatusNotFound)
  829. return
  830. }
  831. failRedirectURL := b.forgetPasswordPageURL
  832. // check reCAPTCHA token
  833. if b.recaptchaEnabled {
  834. token := r.FormValue("token")
  835. if !recaptchaTokenCheck(b, token) {
  836. setFailCodeFlash(w, FailCodeIncorrectRecaptchaToken)
  837. http.Redirect(w, r, failRedirectURL, http.StatusFound)
  838. return
  839. }
  840. }
  841. account := strings.TrimSpace(r.FormValue("account"))
  842. passcode := r.FormValue("otp")
  843. doTOTP := r.URL.Query().Get("totp") == "1"
  844. if doTOTP {
  845. failRedirectURL = MustSetQuery(failRedirectURL, "totp", "1")
  846. }
  847. if account == "" {
  848. setFailCodeFlash(w, FailCodeAccountIsRequired)
  849. http.Redirect(w, r, failRedirectURL, http.StatusFound)
  850. return
  851. }
  852. u, err := b.userModel.(UserPasser).FindUser(b.db, b.newUserObject(), account)
  853. if err != nil {
  854. if err == gorm.ErrRecordNotFound {
  855. setFailCodeFlash(w, FailCodeUserNotFound)
  856. } else {
  857. setFailCodeFlash(w, FailCodeSystemError)
  858. }
  859. setWrongForgetPasswordInputFlash(w, WrongForgetPasswordInputFlash{
  860. Account: account,
  861. })
  862. http.Redirect(w, r, failRedirectURL, http.StatusFound)
  863. return
  864. }
  865. _, createdAt, _ := u.(UserPasser).GetResetPasswordToken()
  866. if createdAt != nil {
  867. v := 60 - int(time.Now().Sub(*createdAt).Seconds())
  868. if v > 0 {
  869. setSecondsToRedoFlash(w, v)
  870. setWrongForgetPasswordInputFlash(w, WrongForgetPasswordInputFlash{
  871. Account: account,
  872. })
  873. http.Redirect(w, r, failRedirectURL, http.StatusFound)
  874. return
  875. }
  876. }
  877. if u.(UserPasser).GetIsTOTPSetup() {
  878. if !doTOTP {
  879. setWrongForgetPasswordInputFlash(w, WrongForgetPasswordInputFlash{
  880. Account: account,
  881. })
  882. failRedirectURL = MustSetQuery(failRedirectURL, "totp", "1")
  883. http.Redirect(w, r, failRedirectURL, http.StatusFound)
  884. return
  885. }
  886. if err = b.consumeTOTPCode(r, u.(UserPasser), passcode); err != nil {
  887. fc := b.getFailCodeFromTOTPCodeConsumeError(err)
  888. setNoticeOrFailCodeFlash(w, err, fc)
  889. setWrongForgetPasswordInputFlash(w, WrongForgetPasswordInputFlash{
  890. Account: account,
  891. TOTP: passcode,
  892. })
  893. http.Redirect(w, r, failRedirectURL, http.StatusFound)
  894. return
  895. }
  896. }
  897. token, err := u.(UserPasser).GenerateResetPasswordToken(b.db, b.newUserObject())
  898. if err != nil {
  899. setFailCodeFlash(w, FailCodeSystemError)
  900. setWrongForgetPasswordInputFlash(w, WrongForgetPasswordInputFlash{
  901. Account: account,
  902. })
  903. http.Redirect(w, r, failRedirectURL, http.StatusFound)
  904. return
  905. }
  906. scheme := "https"
  907. if r.TLS == nil {
  908. scheme = "http"
  909. }
  910. link := fmt.Sprintf("%s://%s%s?id=%s&token=%s", scheme, r.Host, b.resetPasswordPageURL, objectID(u), token)
  911. if doTOTP {
  912. link = MustSetQuery(link, "totp", "1")
  913. }
  914. if b.afterConfirmSendResetPasswordLinkHook != nil {
  915. if herr := b.afterConfirmSendResetPasswordLinkHook(r, u, link); herr != nil {
  916. setNoticeOrFailCodeFlash(w, herr, FailCodeSystemError)
  917. http.Redirect(w, r, failRedirectURL, http.StatusFound)
  918. return
  919. }
  920. }
  921. http.Redirect(w, r, fmt.Sprintf("%s?a=%s", b.resetPasswordLinkSentPageURL, account), http.StatusFound)
  922. return
  923. }
  924. func (b *Builder) doResetPassword(w http.ResponseWriter, r *http.Request) {
  925. if r.Method != http.MethodPost {
  926. w.WriteHeader(http.StatusNotFound)
  927. return
  928. }
  929. userID := r.FormValue("user_id")
  930. token := r.FormValue("token")
  931. passcode := r.FormValue("otp")
  932. doTOTP := r.URL.Query().Get("totp") == "1"
  933. failRedirectURL := fmt.Sprintf("%s?id=%s&token=%s", b.resetPasswordPageURL, userID, token)
  934. if doTOTP {
  935. failRedirectURL = MustSetQuery(failRedirectURL, "totp", "1")
  936. }
  937. if userID == "" {
  938. setFailCodeFlash(w, FailCodeUserNotFound)
  939. http.Redirect(w, r, failRedirectURL, http.StatusFound)
  940. return
  941. }
  942. if token == "" {
  943. setFailCodeFlash(w, FailCodeInvalidToken)
  944. http.Redirect(w, r, failRedirectURL, http.StatusFound)
  945. return
  946. }
  947. password := r.FormValue("password")
  948. confirmPassword := r.FormValue("confirm_password")
  949. if password == "" {
  950. setFailCodeFlash(w, FailCodePasswordCannotBeEmpty)
  951. http.Redirect(w, r, failRedirectURL, http.StatusFound)
  952. return
  953. }
  954. if confirmPassword != password {
  955. setFailCodeFlash(w, FailCodePasswordNotMatch)
  956. setWrongResetPasswordInputFlash(w, WrongResetPasswordInputFlash{
  957. Password: password,
  958. ConfirmPassword: confirmPassword,
  959. })
  960. http.Redirect(w, r, failRedirectURL, http.StatusFound)
  961. return
  962. }
  963. u, err := b.findUserByID(userID)
  964. if err != nil {
  965. if err == ErrUserNotFound {
  966. setFailCodeFlash(w, FailCodeUserNotFound)
  967. } else {
  968. setFailCodeFlash(w, FailCodeSystemError)
  969. }
  970. http.Redirect(w, r, failRedirectURL, http.StatusFound)
  971. return
  972. }
  973. storedToken, _, expired := u.(UserPasser).GetResetPasswordToken()
  974. if expired {
  975. setFailCodeFlash(w, FailCodeTokenExpired)
  976. http.Redirect(w, r, failRedirectURL, http.StatusFound)
  977. return
  978. }
  979. if token != storedToken {
  980. setFailCodeFlash(w, FailCodeInvalidToken)
  981. http.Redirect(w, r, failRedirectURL, http.StatusFound)
  982. return
  983. }
  984. if b.beforeSetPasswordHook != nil {
  985. if herr := b.beforeSetPasswordHook(r, u, password); herr != nil {
  986. setNoticeOrFailCodeFlash(w, herr, FailCodeSystemError)
  987. setWrongResetPasswordInputFlash(w, WrongResetPasswordInputFlash{
  988. Password: password,
  989. ConfirmPassword: confirmPassword,
  990. })
  991. http.Redirect(w, r, failRedirectURL, http.StatusFound)
  992. return
  993. }
  994. }
  995. if u.(UserPasser).GetIsTOTPSetup() {
  996. if !doTOTP {
  997. setWrongResetPasswordInputFlash(w, WrongResetPasswordInputFlash{
  998. Password: password,
  999. ConfirmPassword: confirmPassword,
  1000. })
  1001. failRedirectURL = MustSetQuery(failRedirectURL, "totp", "1")
  1002. http.Redirect(w, r, failRedirectURL, http.StatusFound)
  1003. return
  1004. }
  1005. if err = b.consumeTOTPCode(r, u.(UserPasser), passcode); err != nil {
  1006. fc := b.getFailCodeFromTOTPCodeConsumeError(err)
  1007. setFailCodeFlash(w, fc)
  1008. setWrongResetPasswordInputFlash(w, WrongResetPasswordInputFlash{
  1009. Password: password,
  1010. ConfirmPassword: confirmPassword,
  1011. TOTP: passcode,
  1012. })
  1013. http.Redirect(w, r, failRedirectURL, http.StatusFound)
  1014. return
  1015. }
  1016. }
  1017. err = u.(UserPasser).ConsumeResetPasswordToken(b.db, b.newUserObject())
  1018. if err != nil {
  1019. setFailCodeFlash(w, FailCodeSystemError)
  1020. http.Redirect(w, r, failRedirectURL, http.StatusFound)
  1021. return
  1022. }
  1023. err = u.(UserPasser).SetPassword(b.db, b.newUserObject(), password)
  1024. if err != nil {
  1025. setFailCodeFlash(w, FailCodeSystemError)
  1026. http.Redirect(w, r, failRedirectURL, http.StatusFound)
  1027. return
  1028. }
  1029. if b.afterResetPasswordHook != nil {
  1030. if herr := b.afterResetPasswordHook(r, u); herr != nil {
  1031. setNoticeOrFailCodeFlash(w, herr, FailCodeSystemError)
  1032. http.Redirect(w, r, failRedirectURL, http.StatusFound)
  1033. return
  1034. }
  1035. }
  1036. setInfoCodeFlash(w, InfoCodePasswordSuccessfullyReset)
  1037. http.Redirect(w, r, b.loginPageURL, http.StatusFound)
  1038. return
  1039. }
  1040. // NoticeError
  1041. // ErrWrongPassword
  1042. // ErrEmptyPassword
  1043. // ErrPasswordNotMatch
  1044. // ErrWrongTOTPCode
  1045. // ErrTOTPCodeHasBeenUsed
  1046. func (b *Builder) ChangePassword(
  1047. r *http.Request,
  1048. oldPassword string,
  1049. password string,
  1050. confirmPassword string,
  1051. otp string,
  1052. ) error {
  1053. user := GetCurrentUser(r).(UserPasser)
  1054. if !user.IsPasswordCorrect(oldPassword) {
  1055. return ErrWrongPassword
  1056. }
  1057. if password == "" {
  1058. return ErrEmptyPassword
  1059. }
  1060. if confirmPassword != password {
  1061. return ErrPasswordNotMatch
  1062. }
  1063. if b.beforeSetPasswordHook != nil {
  1064. if herr := b.beforeSetPasswordHook(r, user, password); herr != nil {
  1065. return herr
  1066. }
  1067. }
  1068. if b.totpEnabled {
  1069. if err := b.consumeTOTPCode(r, user, otp); err != nil {
  1070. return err
  1071. }
  1072. }
  1073. err := user.SetPassword(b.db, b.newUserObject(), password)
  1074. if err != nil {
  1075. return err
  1076. }
  1077. if b.afterChangePasswordHook != nil {
  1078. if herr := b.afterChangePasswordHook(r, user); herr != nil {
  1079. return herr
  1080. }
  1081. }
  1082. return nil
  1083. }
  1084. func (b *Builder) doFormChangePassword(w http.ResponseWriter, r *http.Request) {
  1085. if r.Method != http.MethodPost {
  1086. w.WriteHeader(http.StatusNotFound)
  1087. return
  1088. }
  1089. oldPassword := r.FormValue("old_password")
  1090. password := r.FormValue("password")
  1091. confirmPassword := r.FormValue("confirm_password")
  1092. otp := r.FormValue("otp")
  1093. redirectURL := b.changePasswordPageURL
  1094. err := b.ChangePassword(r, oldPassword, password, confirmPassword, otp)
  1095. if err != nil {
  1096. if ne, ok := err.(*NoticeError); ok {
  1097. setNoticeFlash(w, ne)
  1098. } else {
  1099. fc := FailCodeSystemError
  1100. switch err {
  1101. case ErrWrongPassword:
  1102. fc = FailCodeIncorrectPassword
  1103. case ErrEmptyPassword:
  1104. fc = FailCodePasswordCannotBeEmpty
  1105. case ErrPasswordNotMatch:
  1106. fc = FailCodePasswordNotMatch
  1107. case ErrWrongTOTPCode:
  1108. fc = FailCodeIncorrectTOTPCode
  1109. case ErrTOTPCodeHasBeenUsed:
  1110. fc = FailCodeTOTPCodeHasBeenUsed
  1111. }
  1112. setFailCodeFlash(w, fc)
  1113. }
  1114. setWrongChangePasswordInputFlash(w, WrongChangePasswordInputFlash{
  1115. OldPassword: oldPassword,
  1116. NewPassword: password,
  1117. ConfirmPassword: confirmPassword,
  1118. TOTP: otp,
  1119. })
  1120. http.Redirect(w, r, redirectURL, http.StatusFound)
  1121. return
  1122. }
  1123. setInfoCodeFlash(w, InfoCodePasswordSuccessfullyChanged)
  1124. http.Redirect(w, r, b.loginPageURL, http.StatusFound)
  1125. }
  1126. func (b *Builder) totpDo(w http.ResponseWriter, r *http.Request) {
  1127. if r.Method != http.MethodPost {
  1128. w.WriteHeader(http.StatusNotFound)
  1129. return
  1130. }
  1131. var err error
  1132. var user interface{}
  1133. failRedirectURL := b.LogoutURL
  1134. defer func() {
  1135. if err != nil {
  1136. if b.afterFailedToLoginHook != nil {
  1137. if herr := b.afterFailedToLoginHook(r, user, err); herr != nil {
  1138. setNoticeOrFailCodeFlash(w, herr, FailCodeSystemError)
  1139. }
  1140. }
  1141. http.Redirect(w, r, failRedirectURL, http.StatusFound)
  1142. }
  1143. }()
  1144. var claims *UserClaims
  1145. claims, err = parseUserClaimsFromCookie(r, b.authCookieName, b.secret)
  1146. if err != nil {
  1147. return
  1148. }
  1149. user, err = b.findUserByID(claims.UserID)
  1150. if err != nil {
  1151. if err == ErrUserNotFound {
  1152. setFailCodeFlash(w, FailCodeUserNotFound)
  1153. } else {
  1154. setFailCodeFlash(w, FailCodeSystemError)
  1155. }
  1156. return
  1157. }
  1158. u := user.(UserPasser)
  1159. otp := r.FormValue("otp")
  1160. isTOTPSetup := u.GetIsTOTPSetup()
  1161. if err = b.consumeTOTPCode(r, u, otp); err != nil {
  1162. fc := b.getFailCodeFromTOTPCodeConsumeError(err)
  1163. setFailCodeFlash(w, fc)
  1164. failRedirectURL = b.totpValidatePageURL
  1165. if !isTOTPSetup {
  1166. failRedirectURL = b.totpSetupPageURL
  1167. }
  1168. return
  1169. }
  1170. if !isTOTPSetup {
  1171. if err = u.SetIsTOTPSetup(b.db, b.newUserObject(), true); err != nil {
  1172. setFailCodeFlash(w, FailCodeSystemError)
  1173. return
  1174. }
  1175. }
  1176. claims.TOTPValidated = true
  1177. if b.afterLoginHook != nil {
  1178. setCookieForRequest(r, &http.Cookie{Name: b.authCookieName, Value: b.mustGetSessionToken(*claims)})
  1179. if err = b.afterLoginHook(r, user); err != nil {
  1180. setNoticeOrFailCodeFlash(w, err, FailCodeSystemError)
  1181. return
  1182. }
  1183. }
  1184. err = b.setSecureCookiesByClaims(w, user, *claims)
  1185. if err != nil {
  1186. setFailCodeFlash(w, FailCodeSystemError)
  1187. return
  1188. }
  1189. redirectURL := b.homePageURLFunc(r, user)
  1190. if v := b.getContinueURL(w, r); v != "" {
  1191. redirectURL = v
  1192. }
  1193. http.Redirect(w, r, redirectURL, http.StatusFound)
  1194. }
  1195. func (b *Builder) Mount(mux *http.ServeMux) {
  1196. b.MountAPI(mux)
  1197. // pages
  1198. wb := web.New()
  1199. mux.Handle(b.loginPageURL, b.i18nBuilder.EnsureLanguage(wb.Page(b.loginPageFunc)))
  1200. if b.userPassEnabled {
  1201. mux.Handle(b.resetPasswordPageURL, b.i18nBuilder.EnsureLanguage(wb.Page(b.resetPasswordPageFunc)))
  1202. mux.Handle(b.changePasswordPageURL, b.i18nBuilder.EnsureLanguage(wb.Page(b.changePasswordPageFunc)))
  1203. if !b.noForgetPasswordLink {
  1204. mux.Handle(b.forgetPasswordPageURL, b.i18nBuilder.EnsureLanguage(wb.Page(b.forgetPasswordPageFunc)))
  1205. mux.Handle(b.resetPasswordLinkSentPageURL, b.i18nBuilder.EnsureLanguage(wb.Page(b.resetPasswordLinkSentPageFunc)))
  1206. }
  1207. if b.totpEnabled {
  1208. mux.Handle(b.totpSetupPageURL, b.i18nBuilder.EnsureLanguage(wb.Page(b.totpSetupPageFunc)))
  1209. mux.Handle(b.totpValidatePageURL, b.i18nBuilder.EnsureLanguage(wb.Page(b.totpValidatePageFunc)))
  1210. }
  1211. }
  1212. // assets
  1213. assetsSubFS, err := fs.Sub(assetsFS, "assets")
  1214. if err != nil {
  1215. panic(err)
  1216. }
  1217. mux.Handle(assetsPathPrefix, http.StripPrefix(assetsPathPrefix, http.FileServer(http.FS(assetsSubFS))))
  1218. }
  1219. func (b *Builder) MountAPI(mux *http.ServeMux) {
  1220. if len(b.secret) == 0 {
  1221. panic("secret is empty")
  1222. }
  1223. if b.userModel != nil {
  1224. if b.db == nil {
  1225. panic("db is required")
  1226. }
  1227. }
  1228. mux.HandleFunc(b.LogoutURL, b.logout)
  1229. if b.userPassEnabled {
  1230. mux.HandleFunc(b.passwordLoginURL, b.userpassLogin)
  1231. mux.HandleFunc(b.resetPasswordURL, b.doResetPassword)
  1232. mux.HandleFunc(b.changePasswordURL, b.doFormChangePassword)
  1233. if !b.noForgetPasswordLink {
  1234. mux.HandleFunc(b.sendResetPasswordLinkURL, b.sendResetPasswordLink)
  1235. }
  1236. if b.totpEnabled {
  1237. mux.HandleFunc(b.validateTOTPURL, b.totpDo)
  1238. }
  1239. }
  1240. if b.oauthEnabled {
  1241. mux.HandleFunc(b.oauthBeginURL, b.beginAuth)
  1242. mux.HandleFunc(b.oauthCallbackURL, b.completeUserAuthCallback)
  1243. mux.HandleFunc(b.oauthCallbackCompleteURL, b.completeUserAuthCallbackComplete)
  1244. }
  1245. }