From a4a8ceca29e8866fca64dda38f3604c7683c0bb7 Mon Sep 17 00:00:00 2001 From: divyam234 Date: Tue, 22 Aug 2023 20:30:14 +0530 Subject: [PATCH] added 2fa support --- models/file.model.go | 2 +- models/user.model.go | 2 +- routes/auth.go | 2 +- schemas/file.schema.go | 4 +- services/auth.service.go | 87 ++++++++++++++++++++++++++++++---------- services/file.service.go | 6 +-- types/main.go | 2 +- 7 files changed, 74 insertions(+), 31 deletions(-) diff --git a/models/file.model.go b/models/file.model.go index bb6421b..4bd60b4 100644 --- a/models/file.model.go +++ b/models/file.model.go @@ -15,7 +15,7 @@ type File struct { Size int64 `gorm:"type:bigint"` Starred *bool `gorm:"default:false"` Depth *int `gorm:"type:integer"` - UserID int `gorm:"type:bigint;not null"` + UserID int64 `gorm:"type:bigint;not null"` Status string `gorm:"type:text"` ParentID string `gorm:"type:text;index"` Parts *Parts `gorm:"type:jsonb"` diff --git a/models/user.model.go b/models/user.model.go index d39a97d..1ac102c 100644 --- a/models/user.model.go +++ b/models/user.model.go @@ -5,7 +5,7 @@ import ( ) type User struct { - UserId int `gorm:"type:bigint;primaryKey"` + UserId int64 `gorm:"type:bigint;primaryKey"` Name string `gorm:"type:text"` UserName string `gorm:"type:text"` IsPremium bool `gorm:"type:bool"` diff --git a/routes/auth.go b/routes/auth.go index 8914a7c..6263fbd 100644 --- a/routes/auth.go +++ b/routes/auth.go @@ -38,7 +38,7 @@ func addAuthRoutes(rg *gin.RouterGroup) { }) - r.GET("/ws", authService.HandleQrCodeLogin) + r.GET("/ws", authService.HandleMultipleLogin) r.GET("/session", func(c *gin.Context) { diff --git a/schemas/file.schema.go b/schemas/file.schema.go index d9031f2..6056fab 100644 --- a/schemas/file.schema.go +++ b/schemas/file.schema.go @@ -27,7 +27,7 @@ type FileQuery struct { ParentID string `form:"parentId" mapstructure:"parent_id,omitempty"` UpdatedAt *time.Time `form:"updatedAt" mapstructure:"updated_at,omitempty"` Status string `mapstructure:"status"` - UserId int `mapstructure:"user_id"` + UserID int64 `mapstructure:"user_id"` } type FileIn struct { @@ -41,7 +41,7 @@ type FileIn struct { Starred *bool `json:"starred" mapstructure:"starred,omitempty"` Depth *int `json:"depth,omitempty" mapstructure:"depth,omitempty"` Status string `mapstructure:"status,omitempty"` - UserID int `json:"userId" mapstructure:"user_id,omitempty"` + UserID int64 `json:"userId" mapstructure:"user_id,omitempty"` ParentID string `json:"parentId" mapstructure:"parent_id,omitempty"` } diff --git a/services/auth.service.go b/services/auth.service.go index c45553f..f44d48d 100644 --- a/services/auth.service.go +++ b/services/auth.service.go @@ -27,6 +27,7 @@ import ( tgauth "github.com/gotd/td/telegram/auth" "github.com/gotd/td/telegram/auth/qrlogin" "github.com/gotd/td/tg" + "github.com/gotd/td/tgerr" "gorm.io/gorm" ) @@ -46,6 +47,7 @@ type SocketMessage struct { PhoneNo string `json:"phoneNo,omitempty"` PhoneCodeHash string `json:"phoneCodeHash,omitempty"` PhoneCode string `json:"phoneCode,omitempty"` + Password string `json:"password,omitempty"` } func IP4toInt(IPv4Address net.IP) int64 { @@ -109,7 +111,7 @@ func (as *AuthService) LogIn(c *gin.Context) (*schemas.Message, *types.AppError) now := time.Now().UTC() jwtClaims := &types.JWTClaims{Claims: jwt.Claims{ - Subject: strconv.Itoa(session.UserID), + Subject: strconv.FormatInt(session.UserID, 10), IssuedAt: jwt.NewNumericDate(now), Expiry: jwt.NewNumericDate(now.Add(time.Duration(as.SessionMaxAge) * time.Second)), }, TgSession: session.Sesssion, @@ -219,7 +221,7 @@ func prepareSession(user *tg.User, data *session.Data) *types.TgSession { sessionString := generateTgSession(data.DC, data.AuthKey, 443) session := &types.TgSession{ Sesssion: sessionString, - UserID: int(user.ID), + UserID: user.ID, Bot: user.Bot, UserName: user.Username, Name: fmt.Sprintf("%s %s", user.FirstName, user.LastName), @@ -228,7 +230,7 @@ func prepareSession(user *tg.User, data *session.Data) *types.TgSession { return session } -func (as *AuthService) HandleQrCodeLogin(c *gin.Context) { +func (as *AuthService) HandleMultipleLogin(c *gin.Context) { upgrader := websocket.Upgrader{ CheckOrigin: func(r *http.Request) bool { return true @@ -255,11 +257,19 @@ func (as *AuthService) HandleQrCodeLogin(c *gin.Context) { conn.WriteJSON(map[string]interface{}{"type": "auth", "payload": map[string]string{"token": token.URL()}}) return nil }) + + if tgerr.Is(err, "SESSION_PASSWORD_NEEDED") { + conn.WriteJSON(map[string]interface{}{"type": "auth", "message": "2FA required"}) + return + } + if err != nil { + conn.WriteJSON(map[string]interface{}{"type": "error", "message": err.Error()}) return } user, ok := authorization.User.AsNotEmpty() if !ok { + conn.WriteJSON(map[string]interface{}{"type": "error", "message": errors.New("auth failed")}) return } res, _ := sessionStorage.LoadSession(c) @@ -270,27 +280,60 @@ func (as *AuthService) HandleQrCodeLogin(c *gin.Context) { }() } if message.AuthType == "phone" && message.Message == "sendcode" { - res, err := tgClient.Auth().SendCode(c, message.PhoneNo, tgauth.SendCodeOptions{}) - if err != nil { - return - } - code := res.(*tg.AuthSentCode) - conn.WriteJSON(map[string]interface{}{"type": "auth", "payload": map[string]string{"phoneCodeHash": code.PhoneCodeHash}}) + go func() { + res, err := tgClient.Auth().SendCode(c, message.PhoneNo, tgauth.SendCodeOptions{}) + if err != nil { + conn.WriteJSON(map[string]interface{}{"type": "error", "message": err.Error()}) + return + } + code := res.(*tg.AuthSentCode) + conn.WriteJSON(map[string]interface{}{"type": "auth", "payload": map[string]string{"phoneCodeHash": code.PhoneCodeHash}}) + }() } if message.AuthType == "phone" && message.Message == "signin" { - auth, err := tgClient.Auth().SignIn(c, message.PhoneNo, message.PhoneCode, message.PhoneCodeHash) - if err != nil { - return - } - user, ok := auth.User.AsNotEmpty() - if !ok { - return - } - res, _ := sessionStorage.LoadSession(c) - sessionData := &SessionData{} - json.Unmarshal(res, sessionData) - session := prepareSession(user, &sessionData.Data) - conn.WriteJSON(map[string]interface{}{"type": "auth", "payload": session, "message": "success"}) + go func() { + auth, err := tgClient.Auth().SignIn(c, message.PhoneNo, message.PhoneCode, message.PhoneCodeHash) + + if errors.Is(err, tgauth.ErrPasswordAuthNeeded) { + conn.WriteJSON(map[string]interface{}{"type": "auth", "message": "2FA required"}) + return + } + + if err != nil { + conn.WriteJSON(map[string]interface{}{"type": "error", "message": err.Error()}) + return + } + user, ok := auth.User.AsNotEmpty() + if !ok { + conn.WriteJSON(map[string]interface{}{"type": "error", "message": errors.New("auth failed")}) + return + } + res, _ := sessionStorage.LoadSession(c) + sessionData := &SessionData{} + json.Unmarshal(res, sessionData) + session := prepareSession(user, &sessionData.Data) + conn.WriteJSON(map[string]interface{}{"type": "auth", "payload": session, "message": "success"}) + }() + } + + if message.AuthType == "2fa" && message.Password != "" { + go func() { + auth, err := tgClient.Auth().Password(c, message.Password) + if err != nil { + conn.WriteJSON(map[string]interface{}{"type": "error", "message": err.Error()}) + return + } + user, ok := auth.User.AsNotEmpty() + if !ok { + conn.WriteJSON(map[string]interface{}{"type": "error", "message": errors.New("auth failed")}) + return + } + res, _ := sessionStorage.LoadSession(c) + sessionData := &SessionData{} + json.Unmarshal(res, sessionData) + session := prepareSession(user, &sessionData.Data) + conn.WriteJSON(map[string]interface{}{"type": "auth", "payload": session, "message": "success"}) + }() } if err != nil { log.Println(err) diff --git a/services/file.service.go b/services/file.service.go index afe6667..d1215f1 100644 --- a/services/file.service.go +++ b/services/file.service.go @@ -33,10 +33,10 @@ type FileService struct { ChannelID int64 } -func getAuthUserId(c *gin.Context) int { +func getAuthUserId(c *gin.Context) int64 { val, _ := c.Get("jwtUser") jwtUser := val.(*types.JWTClaims) - userId, _ := strconv.Atoi(jwtUser.Subject) + userId, _ := strconv.ParseInt(jwtUser.Subject, 10, 64) return userId } @@ -160,7 +160,7 @@ func (fs *FileService) ListFiles(c *gin.Context) (*schemas.FileResponse, *types. var fileQuery schemas.FileQuery fileQuery.Op = "list" fileQuery.Status = "active" - fileQuery.UserId = userId + fileQuery.UserID = userId if err := c.ShouldBindQuery(&fileQuery); err != nil { return nil, &types.AppError{Error: errors.New(""), Code: http.StatusBadRequest} } diff --git a/types/main.go b/types/main.go index 6af9ecf..5ed430c 100644 --- a/types/main.go +++ b/types/main.go @@ -28,7 +28,7 @@ type JWTClaims struct { type TgSession struct { Sesssion string `json:"session"` - UserID int `json:"userId"` + UserID int64 `json:"userId"` Bot bool `json:"bot"` UserName string `json:"userName"` Name string `json:"name"`