Allow restricting file download

server/fileservice.go

@@ -3,11 +3,15 @@ package server
 import (
 	"context"
 	"fmt"
+	"time"
 
 	"gitea.benny.dog/torjus/ezshare/pb"
 	"gitea.benny.dog/torjus/ezshare/store"
 	"github.com/dustin/go-humanize"
+	"github.com/google/uuid"
 	"go.uber.org/zap"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 )
 
 type GRPCFileServiceServer struct {
@@ -24,14 +28,27 @@ func NewGRPCFileServiceServer(store store.FileStore) *GRPCFileServiceServer {
 func (s *GRPCFileServiceServer) UploadFile(ctx context.Context, req *pb.UploadFileRequest) (*pb.UploadFileResponse, error) {
 	var f pb.File
 	f.Data = req.GetData()
+	f.Metadata = &pb.File_Metadata{}
+	f.Metadata.OriginalFilename = req.OriginalFilename
+	if req.WithPasscode {
+		f.Metadata.Passcode = uuid.Must(uuid.NewRandom()).String()
+	}
+
+	f.Metadata.ExpiresOn = req.ExpiresOn
+
+	f.Metadata.MaxViews = req.MaxViews
 
 	id, err := s.store.StoreFile(&f)
 	if err != nil {
 		s.Logger.Warnw("Error storing file.", "error", err)
 		return nil, err
 	}
+	fileUrl := fmt.Sprintf("%s/files/%s", s.Hostname, id)
+	if req.WithPasscode {
+		fileUrl = fmt.Sprintf("%s?passcode=%s", fileUrl, f.Metadata.Passcode)
+	}
 	s.Logger.Infow("Received file.", "id", id, "size", humanize.Bytes(uint64(len(f.Data))))
-	return &pb.UploadFileResponse{Id: id, FileUrl: fmt.Sprintf("%s/files/%s", s.Hostname, id)}, nil
+	return &pb.UploadFileResponse{Id: id, FileUrl: fileUrl}, nil
 }
 
 func (s *GRPCFileServiceServer) GetFile(ctx context.Context, req *pb.GetFileRequest) (*pb.GetFileResponse, error) {
@@ -39,8 +56,40 @@ func (s *GRPCFileServiceServer) GetFile(ctx context.Context, req *pb.GetFileRequ
 	if err != nil {
 		return nil, err
 	}
-	s.Logger.Infow("Sending file to client.", "id", f.FileId, "size", humanize.Bytes(uint64(len(f.Data))))
 
+	// Check if expired, viewed too many times, or needs passcode
+	if f.Metadata.ExpiresOn != nil {
+		if f.Metadata.ExpiresOn.AsTime().Before(time.Now()) {
+			if err := s.store.DeleteFile(f.FileId); err != nil {
+				s.Logger.Warnw("Error deleting expired file.", "error", err)
+			}
+			return nil, status.Error(codes.NotFound, "no such file")
+		}
+	}
+
+	// Check if requires passcode
+	if f.Metadata.Passcode != "" {
+		if f.Metadata.Passcode != req.Passcode {
+			return nil, status.Error(codes.PermissionDenied, "needs passcode")
+		}
+	}
+
+	// Too many views
+	if f.Metadata.MaxViews != 0 {
+		if f.Metadata.CurrentViews >= f.Metadata.MaxViews {
+			if err := s.store.DeleteFile(f.FileId); err != nil {
+				s.Logger.Warnw("Error deleting expired file.", "error", err)
+			}
+			return nil, status.Error(codes.NotFound, "no such file")
+		}
+	}
+
+	f.Metadata.CurrentViews++
+	if _, err := s.store.StoreFile(f); err != nil {
+		s.Logger.Warnw("Error storing updated file.", "error", err)
+	}
+
+	s.Logger.Infow("Sending file to client.", "id", f.FileId, "size", humanize.Bytes(uint64(len(f.Data))))
 	return &pb.GetFileResponse{File: f}, nil
 }
 

server/http.go

@@ -4,6 +4,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"net/http"
+	"time"
 
 	"gitea.benny.dog/torjus/ezshare/store"
 	"github.com/go-chi/chi/v5"
@@ -72,11 +73,51 @@ func (s *HTTPServer) FileHandler(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
+	// Check if expired, passcode required, or other reasons for not sending file
+	// To many views
+	if f.Metadata.MaxViews != 0 {
+		if f.Metadata.CurrentViews >= f.Metadata.MaxViews {
+			// File has been viewed max amount of times - delete
+			s.Logger.Infow("File has been viewed too many times.", "id", id)
+			if err := s.store.DeleteFile(id); err != nil {
+				s.Logger.Warnw("Error deleting file.", "id", id)
+			}
+			WriteErrorResponse(w, http.StatusNotFound, "file viewed too many times")
+			return
+		}
+	}
+
+	// Expired
+	if f.Metadata.ExpiresOn != nil {
+		if f.Metadata.ExpiresOn.AsTime().Before(time.Now()) {
+			// File has expired
+			s.Logger.Infow("File has been expired.", "id", id)
+			if err := s.store.DeleteFile(id); err != nil {
+				s.Logger.Warnw("Error deleting file.", "id", id)
+			}
+			WriteErrorResponse(w, http.StatusNotFound, "file not found")
+			return
+		}
+	}
+
+	// Passcode required
+	if f.Metadata.Passcode != "" {
+		if f.Metadata.Passcode != r.URL.Query().Get("passcode") {
+			WriteErrorResponse(w, http.StatusUnauthorized, "wrong passcode")
+			return
+		}
+	}
+
 	w.Header().Add("Content-Type", http.DetectContentType(f.Data))
 	if _, err := w.Write(f.Data); err != nil {
 		s.Logger.Warnw("Error sending file.", "error", err, "remote_addr", r.RemoteAddr)
 	}
 	s.Logger.Infow("Sent file.", "remote_addr", r.RemoteAddr)
+	f.Metadata.CurrentViews++
+
+	if _, err := s.store.StoreFile(f); err != nil {
+		s.Logger.Warnw("Error storing updated file.", "error", err)
+	}
 }
 
 func WriteErrorResponse(w http.ResponseWriter, status int, message string) {