added 2fa support

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"`

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"`

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) {
 

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"`
 }
 

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{})
+			go func() {
-			if err != nil {
+				res, err := tgClient.Auth().SendCode(c, message.PhoneNo, tgauth.SendCodeOptions{})
-				return
+				if err != nil {
-			}
+					conn.WriteJSON(map[string]interface{}{"type": "error", "message": err.Error()})
-			code := res.(*tg.AuthSentCode)
+					return
-			conn.WriteJSON(map[string]interface{}{"type": "auth", "payload": map[string]string{"phoneCodeHash": code.PhoneCodeHash}})
+				}
+				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)
+			go func() {
-			if err != nil {
+				auth, err := tgClient.Auth().SignIn(c, message.PhoneNo, message.PhoneCode, message.PhoneCodeHash)
-				return
+
-			}
+				if errors.Is(err, tgauth.ErrPasswordAuthNeeded) {
-			user, ok := auth.User.AsNotEmpty()
+					conn.WriteJSON(map[string]interface{}{"type": "auth", "message": "2FA required"})
-			if !ok {
+					return
-				return
+				}
-			}
+
-			res, _ := sessionStorage.LoadSession(c)
+				if err != nil {
-			sessionData := &SessionData{}
+					conn.WriteJSON(map[string]interface{}{"type": "error", "message": err.Error()})
-			json.Unmarshal(res, sessionData)
+					return
-			session := prepareSession(user, &sessionData.Data)
+				}
-			conn.WriteJSON(map[string]interface{}{"type": "auth", "payload": session, "message": "success"})
+				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)

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}
 	}

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"`