From fdd6e64c3f19167568b64ea91d58fea7f51db40c Mon Sep 17 00:00:00 2001
From: divyam234 <47589864+divyam234@users.noreply.github.com>
Date: Wed, 1 Jan 2025 01:17:51 +0530
Subject: [PATCH] feat: enhance authentication and file sharing features

---
 internal/auth/auth.go              |  5 +++++
 pkg/services/auth.go               |  2 +-
 pkg/services/file.go               | 11 ++++++++++-
 pkg/services/file_query_builder.go |  3 ++-
 pkg/services/user.go               |  6 ++++--
 5 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/internal/auth/auth.go b/internal/auth/auth.go
index a99838c..f9bd861 100644
--- a/internal/auth/auth.go
+++ b/internal/auth/auth.go
@@ -48,6 +48,11 @@ func GetUser(c context.Context) (int64, string) {
 	return userId, authUser.TgSession
 }
 
+func GetJWTUser(c context.Context) *types.JWTClaims {
+	authUser, _ := c.Value(authKey).(*types.JWTClaims)
+	return authUser
+}
+
 func VerifyUser(db *gorm.DB, cache cache.Cacher, secret, authCookie string) (*types.JWTClaims, error) {
 	claims, err := Decode(secret, authCookie)
 
diff --git a/pkg/services/auth.go b/pkg/services/auth.go
index 0f7f018..8d4deb4 100644
--- a/pkg/services/auth.go
+++ b/pkg/services/auth.go
@@ -121,7 +121,7 @@ func (a *apiService) AuthLogin(ctx context.Context, session *api.SessionCreate)
 }
 
 func (a *apiService) AuthLogout(ctx context.Context) (*api.AuthLogoutNoContent, error) {
-	authUser, _ := ctx.Value("authUser").(*types.JWTClaims)
+	authUser := auth.GetJWTUser(ctx)
 	client, _ := tgc.AuthClient(ctx, &a.cnf.TG, authUser.TgSession)
 	tgc.RunWithAuth(ctx, client, "", func(ctx context.Context) error {
 		_, err := client.API().AuthLogOut(ctx)
diff --git a/pkg/services/file.go b/pkg/services/file.go
index 3a5c178..1a55adb 100644
--- a/pkg/services/file.go
+++ b/pkg/services/file.go
@@ -443,7 +443,16 @@ func (a *apiService) FilesShareByid(ctx context.Context, params api.FilesShareBy
 	if len(result) == 0 {
 		return nil, notFoundErr
 	}
-	return &api.FileShare{ExpiresAt: api.NewOptDateTime(*result[0].ExpiresAt), Protected: result[0].Password != nil}, nil
+	res := &api.FileShare{
+		ID: result[0].ID,
+	}
+	if result[0].Password != nil {
+		res.Protected = true
+	}
+	if result[0].ExpiresAt != nil {
+		res.ExpiresAt = api.NewOptDateTime(*result[0].ExpiresAt)
+	}
+	return res, nil
 }
 
 func (a *apiService) FilesStream(ctx context.Context, params api.FilesStreamParams) (api.FilesStreamRes, error) {
diff --git a/pkg/services/file_query_builder.go b/pkg/services/file_query_builder.go
index b3ed93c..e4499e7 100644
--- a/pkg/services/file_query_builder.go
+++ b/pkg/services/file_query_builder.go
@@ -1,6 +1,7 @@
 package services
 
 import (
+	"errors"
 	"fmt"
 	"math"
 	"strings"
@@ -39,7 +40,7 @@ func (afb *fileQueryBuilder) execute(filesQuery *api.FilesListParams, userId int
 	res := []fileResponse{}
 	if err := query.Scan(&res).Error; err != nil {
 		if strings.Contains(err.Error(), "file not found") {
-			return nil, &apiError{err: err}
+			return nil, &apiError{err: errors.New("invalid path"), code: 404}
 		}
 		return nil, &apiError{err: err}
 	}
diff --git a/pkg/services/user.go b/pkg/services/user.go
index 47b5e65..9d525f6 100644
--- a/pkg/services/user.go
+++ b/pkg/services/user.go
@@ -132,7 +132,7 @@ func (a *apiService) UsersListSessions(ctx context.Context) ([]api.UserSession,
 	return sessionsOut, nil
 }
 
-func (a *apiService) UsersProfileImage(ctx context.Context) (*api.UsersProfileImageOKHeaders, error) {
+func (a *apiService) UsersProfileImage(ctx context.Context, params api.UsersProfileImageParams) (*api.UsersProfileImageOKHeaders, error) {
 	_, session := auth.GetUser(ctx)
 
 	client, err := tgc.AuthClient(ctx, &a.cnf.TG, session)
@@ -156,14 +156,16 @@ func (a *apiService) UsersProfileImage(ctx context.Context) (*api.UsersProfileIm
 		if !ok {
 			return errors.New("profile not found")
 		}
+		photo.GetPersonal()
 		location := &tg.InputPeerPhotoFileLocation{Big: false, Peer: peer, PhotoID: photo.PhotoID}
 		buff, err := tgc.GetMediaContent(ctx, client.API(), location)
 		if err != nil {
 			return err
 		}
 		content := buff.Bytes()
-		res.SetCacheControl("public, max-age=86400")
+		res.SetCacheControl("public, max-age=86400, must-revalidate")
 		res.SetContentLength(int64(len(content)))
+		res.SetEtag(fmt.Sprintf("\"%v\"", photo.PhotoID))
 		res.SetContentDisposition(fmt.Sprintf("inline; filename=\"%s\"", "profile.jpeg"))
 		res.Response = api.UsersProfileImageOK{Data: bytes.NewReader(content)}
 		return nil