diff --git a/go.mod b/go.mod index 96514b8..28e53a4 100644 --- a/go.mod +++ b/go.mod @@ -36,6 +36,7 @@ require ( github.com/go-playground/universal-translator v0.18.1 // indirect github.com/go-playground/validator/v10 v10.14.0 // indirect github.com/goccy/go-json v0.10.2 // indirect + github.com/gorilla/websocket v1.5.0 github.com/gotd/ige v0.2.2 // indirect github.com/gotd/neo v0.1.5 // indirect github.com/jackc/pgpassfile v1.0.0 // indirect diff --git a/go.sum b/go.sum index b38bac4..dc8d265 100644 --- a/go.sum +++ b/go.sum @@ -64,6 +64,8 @@ github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeN github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= github.com/gorilla/websocket v1.4.1 h1:q7AeDBpnBk8AogcD4DSag/Ukw/KV+YhzLj2bP5HvKCM= github.com/gorilla/websocket v1.4.1/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE= +github.com/gorilla/websocket v1.5.0 h1:PPwGk2jz7EePpoHN/+ClbZu8SPxiqlu12wZP/3sWmnc= +github.com/gorilla/websocket v1.5.0/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE= github.com/gotd/contrib v0.19.0 h1:O6GvMrRVeFslIHLUcpaHVzcl9/5PcgR2jQTIIeTyds0= github.com/gotd/contrib v0.19.0/go.mod h1:LzPxzRF0FvtpBt/WyODWQnPpk0tm/G9z6RHUoPqMakU= github.com/gotd/ige v0.2.2 h1:XQ9dJZwBfDnOGSTxKXBGP4gMud3Qku2ekScRjDWWfEk= diff --git a/routes/auth.go b/routes/auth.go index f0b10d1..f981c7b 100644 --- a/routes/auth.go +++ b/routes/auth.go @@ -38,6 +38,8 @@ func addAuthRoutes(rg *gin.RouterGroup) { }) + r.GET("/ws", authService.HandleQrCodeLogin) + r.GET("/session", func(c *gin.Context) { session := authService.GetSession(c) diff --git a/services/auth.service.go b/services/auth.service.go index 3fc28ff..6feb53f 100644 --- a/services/auth.service.go +++ b/services/auth.service.go @@ -5,8 +5,10 @@ import ( "context" "encoding/base64" "encoding/binary" - "encoding/hex" + "encoding/json" "errors" + "fmt" + "log" "math/big" "net" "net/http" @@ -20,6 +22,11 @@ import ( "github.com/divyam234/teldrive-go/utils/auth" "github.com/gin-gonic/gin" "github.com/go-jose/go-jose/v3/jwt" + "github.com/gorilla/websocket" + "github.com/gotd/td/session" + tgauth "github.com/gotd/td/telegram/auth" + "github.com/gotd/td/telegram/auth/qrlogin" + "github.com/gotd/td/tg" "gorm.io/gorm" "gorm.io/gorm/clause" ) @@ -29,6 +36,18 @@ type AuthService struct { SessionMaxAge int } +type SessionData struct { + Version int + Data session.Data +} +type SocketMessage struct { + AuthType string `json:"authType"` + Message string `json:"message"` + PhoneNo string `json:"phoneNo,omitempty"` + PhoneCodeHash string `json:"phoneCodeHash,omitempty"` + PhoneCode string `json:"phoneCode,omitempty"` +} + func IP4toInt(IPv4Address net.IP) int64 { IPv4Int := big.NewInt(0) IPv4Int.SetBytes(IPv4Address.To4()) @@ -43,10 +62,18 @@ func Pack32BinaryIP4(ip4Address string) []byte { return buf.Bytes() } -func generateTgSession(dcID int, serverAddress string, authKey []byte, port int) string { +func generateTgSession(dcID int, authKey []byte, port int) string { + + dcMaps := map[int]string{ + 1: "149.154.175.53", + 2: "149.154.167.51", + 3: "149.154.175.100", + 4: "149.154.167.91", + 5: "91.108.56.130", + } dcIDByte := byte(dcID) - serverAddressBytes := Pack32BinaryIP4(serverAddress) + serverAddressBytes := Pack32BinaryIP4(dcMaps[dcID]) portByte := make([]byte, 2) binary.BigEndian.PutUint16(portByte, uint16(port)) @@ -74,13 +101,6 @@ func GetUserSessionCookieName(c *gin.Context) string { } func (as *AuthService) LogIn(c *gin.Context) (*schemas.Message, *types.AppError) { - dcMaps := map[int]string{ - 1: "149.154.175.53", - 2: "149.154.167.51", - 3: "149.154.175.100", - 4: "149.154.167.91", - 5: "91.108.56.130", - } var session types.TgSession if err := c.ShouldBindJSON(&session); err != nil { return nil, &types.AppError{Error: errors.New("invalid request payload"), Code: http.StatusBadRequest} @@ -88,15 +108,16 @@ func (as *AuthService) LogIn(c *gin.Context) (*schemas.Message, *types.AppError) now := time.Now().UTC() - authBytes, _ := hex.DecodeString(session.AuthKey) - - sessionData := generateTgSession(session.DcID, dcMaps[session.DcID], authBytes, 443) - jwtClaims := &types.JWTClaims{Claims: jwt.Claims{ - Subject: session.UserID, + Subject: strconv.Itoa(session.UserID), IssuedAt: jwt.NewNumericDate(now), Expiry: jwt.NewNumericDate(now.Add(time.Duration(as.SessionMaxAge) * time.Second)), - }, TgSession: sessionData, Name: session.Name, UserName: session.UserName, Bot: session.Bot, IsPremium: session.IsPremium} + }, TgSession: session.Sesssion, + Name: session.Name, + UserName: session.UserName, + Bot: session.Bot, + IsPremium: session.IsPremium, + } jweToken, err := auth.Encode(jwtClaims) @@ -104,17 +125,16 @@ func (as *AuthService) LogIn(c *gin.Context) (*schemas.Message, *types.AppError) return nil, &types.AppError{Error: err, Code: http.StatusBadRequest} } - userId, _ := strconv.Atoi(session.UserID) user := models.User{ - UserId: userId, + UserId: session.UserID, Name: session.Name, UserName: session.UserName, IsPremium: session.IsPremium, - TgSession: sessionData, + TgSession: session.Sesssion, } if err := as.Db.Clauses(clause.OnConflict{ Columns: []clause.Column{{Name: "user_id"}}, - DoUpdates: clause.Assignments(map[string]interface{}{"tg_session": sessionData}), + DoUpdates: clause.Assignments(map[string]interface{}{"tg_session": session.Sesssion}), }).Create(&user).Error; err != nil { return nil, &types.AppError{Error: errors.New("failed to create or update user"), Code: http.StatusInternalServerError} } @@ -166,11 +186,90 @@ func (as *AuthService) Logout(c *gin.Context) (*schemas.Message, *types.AppError return nil, &types.AppError{Error: err, Code: http.StatusInternalServerError} } - _, err = tgClient.Tg.API().AuthLogOut(context.Background()) - if err != nil { - return nil, &types.AppError{Error: err, Code: http.StatusInternalServerError} - } + tgClient.Tg.API().AuthLogOut(c) utils.StopClient(stop, userId) c.SetCookie(GetUserSessionCookieName(c), "", -1, "/", c.Request.Host, false, false) return &schemas.Message{Status: true, Message: "logout success"}, nil } + +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), + Bot: user.Bot, + UserName: user.Username, + Name: fmt.Sprintf("%s %s", user.FirstName, user.LastName), + IsPremium: user.Premium, + } + return session +} + +func (as *AuthService) HandleQrCodeLogin(c *gin.Context) { + upgrader := websocket.Upgrader{ + CheckOrigin: func(r *http.Request) bool { + return true + }, + } + conn, err := upgrader.Upgrade(c.Writer, c.Request, nil) + if err != nil { + log.Println(err) + return + } + defer conn.Close() + + dispatcher := tg.NewUpdateDispatcher() + loggedIn := qrlogin.OnLoginToken(dispatcher) + sessionStorage := &session.StorageMemory{} + tgClient, stop, _ := utils.GetNonAuthClient(dispatcher, sessionStorage) + defer stop() + for { + message := &SocketMessage{} + err := conn.ReadJSON(message) + if message.AuthType == "qr" { + authorization, err := tgClient.QR().Auth(c, loggedIn, func(ctx context.Context, token qrlogin.Token) error { + conn.WriteJSON(map[string]interface{}{"type": "auth", "payload": map[string]string{"token": token.URL()}}) + return nil + }) + if err != nil { + break + } + user, ok := authorization.User.AsNotEmpty() + if !ok { + break + } + 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 == "phone" && message.Message == "sendcode" { + res, err := tgClient.Auth().SendCode(c, message.PhoneNo, tgauth.SendCodeOptions{}) + if err != nil { + break + } + 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 { + break + } + user, ok := auth.User.AsNotEmpty() + if !ok { + break + } + 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) + return + } + } +} diff --git a/types/main.go b/types/main.go index 6e2f0d8..6af9ecf 100644 --- a/types/main.go +++ b/types/main.go @@ -27,9 +27,8 @@ type JWTClaims struct { } type TgSession struct { - DcID int `json:"dcId"` - AuthKey string `json:"authKey"` - UserID string `json:"userId"` + Sesssion string `json:"session"` + UserID int `json:"userId"` Bot bool `json:"bot"` UserName string `json:"userName"` Name string `json:"name"` diff --git a/utils/tgclient.go b/utils/tgclient.go index b88fa26..92cd201 100644 --- a/utils/tgclient.go +++ b/utils/tgclient.go @@ -166,6 +166,25 @@ func GetBotClient() *Client { return smallest } +func GetNonAuthClient(handler telegram.UpdateHandler, storage telegram.SessionStorage) (*telegram.Client, bg.StopFunc, error) { + + client := telegram.NewClient(config.AppId, config.AppHash, telegram.Options{ + SessionStorage: storage, + Middlewares: []telegram.Middleware{ + ratelimit.New(rate.Every(time.Millisecond*100), 5), + }, + UpdateHandler: handler, + }) + + stop, err := bg.Connect(client) + + if err != nil { + return nil, nil, err + } + + return client, stop, nil +} + func StopClient(stop bg.StopFunc, key int) { stop() delete(clients, key)