From 35b49739ec7a79e2b2fe9080007e7d885606cb10 Mon Sep 17 00:00:00 2001 From: Josh Yan Date: Wed, 17 Jul 2024 14:53:29 -0700 Subject: [PATCH] timecheck --- cmd/cmd.go | 1 + server/routes.go | 10 ++++++++++ 2 files changed, 11 insertions(+) diff --git a/cmd/cmd.go b/cmd/cmd.go index 06bc13332..768b5c841 100644 --- a/cmd/cmd.go +++ b/cmd/cmd.go @@ -343,6 +343,7 @@ func getLocalPath(ctx context.Context, digest string) (string, error) { } request.Header.Set("Authorization", authz) + request.Header.Set("Timestamp", time.Now().Format(time.RFC3339)) request.Header.Set("User-Agent", fmt.Sprintf("ollama/%s (%s %s) Go/%s", version.Version, runtime.GOARCH, runtime.GOOS, runtime.Version())) request.Header.Set("X-Redirect-Create", "1") diff --git a/server/routes.go b/server/routes.go index 7f55a36cf..31328dc46 100644 --- a/server/routes.go +++ b/server/routes.go @@ -997,6 +997,16 @@ func (s *Server) IsServerKeyPublicKey(c *gin.Context) bool { serverPublicKey, err := auth.GetPublicKey() if err != nil { slog.Error(fmt.Sprintf("failed to get server public key: %v", err)) + return false + } + + timestamp, err := time.Parse(time.RFC3339, c.GetHeader("Timestamp")) + if err != nil { + return false + } + + if time.Since(timestamp) > time.Minute { + return false } if bytes.Equal(serverPublicKey.Marshal(), clientPublicKey.Marshal()) {