Fixed authentication key generation for partition key ranges

Fixed collection rid generation

Improved compatibility with SDKs

api/handlers/collections.go

@@ -1,6 +1,7 @@
 package handlers
 
 import (
+	"fmt"
 	"net/http"
 
 	"github.com/gin-gonic/gin"
@@ -15,6 +16,7 @@ func GetAllCollections(c *gin.Context) {
 	if status == repositorymodels.StatusOk {
 		database, _ := repositories.GetDatabase(databaseId)
 
+		c.Header("x-ms-item-count", fmt.Sprintf("%d", len(collections)))
 		c.IndentedJSON(http.StatusOK, gin.H{
 			"_rid":                database.ResourceID,
 			"DocumentCollections": collections,

api/handlers/databases.go

@@ -1,6 +1,7 @@
 package handlers
 
 import (
+	"fmt"
 	"net/http"
 
 	"github.com/gin-gonic/gin"
@@ -11,6 +12,7 @@ import (
 func GetAllDatabases(c *gin.Context) {
 	databases, status := repositories.GetAllDatabases()
 	if status == repositorymodels.StatusOk {
+		c.Header("x-ms-item-count", fmt.Sprintf("%d", len(databases)))
 		c.IndentedJSON(http.StatusOK, gin.H{
 			"_rid":      "",
 			"Databases": databases,

api/handlers/documents.go

@@ -1,6 +1,7 @@
 package handlers
 
 import (
+	"fmt"
 	"net/http"
 
 	"github.com/gin-gonic/gin"
@@ -17,6 +18,7 @@ func GetAllDocuments(c *gin.Context) {
 	if status == repositorymodels.StatusOk {
 		collection, _ := repositories.GetCollection(databaseId, collectionId)
 
+		c.Header("x-ms-item-count", fmt.Sprintf("%d", len(documents)))
 		c.IndentedJSON(http.StatusOK, gin.H{
 			"_rid":      collection.ID,
 			"Documents": documents,
@@ -128,6 +130,7 @@ func DocumentsPost(c *gin.Context) {
 		}
 
 		collection, _ := repositories.GetCollection(databaseId, collectionId)
+		c.Header("x-ms-item-count", fmt.Sprintf("%d", len(docs)))
 		c.IndentedJSON(http.StatusOK, gin.H{
 			"_rid":      collection.ResourceID,
 			"Documents": docs,

api/handlers/middleware/authentication.go

@@ -44,6 +44,11 @@ func Authentication() gin.HandlerFunc {
 			resourceId += "/docs/" + docId
 		}
 
+		isFeed := c.Request.Header.Get("A-Im") == "Incremental Feed"
+		if resourceType == "pkranges" && isFeed {
+			resourceId = collId
+		}
+
 		authHeader := c.Request.Header.Get("authorization")
 		date := c.Request.Header.Get("x-ms-date")
 		expectedSignature := authentication.GenerateSignature(

api/handlers/offers.go

@@ -7,6 +7,7 @@ import (
 )
 
 func GetOffers(c *gin.Context) {
+	c.Header("x-ms-item-count", "0")
 	c.IndentedJSON(http.StatusOK, gin.H{
 		"_rid":   "",
 		"_count": 0,

api/handlers/stored_procedures.go

@@ -1,6 +1,7 @@
 package handlers
 
 import (
+	"fmt"
 	"net/http"
 
 	"github.com/gin-gonic/gin"
@@ -15,6 +16,7 @@ func GetAllStoredProcedures(c *gin.Context) {
 	sps, status := repositories.GetAllStoredProcedures(databaseId, collectionId)
 
 	if status == repositorymodels.StatusOk {
+		c.Header("x-ms-item-count", fmt.Sprintf("%d", len(sps)))
 		c.IndentedJSON(http.StatusOK, gin.H{"_rid": "", "StoredProcedures": sps, "_count": len(sps)})
 		return
 	}

api/handlers/triggers.go

@@ -1,6 +1,7 @@
 package handlers
 
 import (
+	"fmt"
 	"net/http"
 
 	"github.com/gin-gonic/gin"
@@ -15,6 +16,7 @@ func GetAllTriggers(c *gin.Context) {
 	triggers, status := repositories.GetAllTriggers(databaseId, collectionId)
 
 	if status == repositorymodels.StatusOk {
+		c.Header("x-ms-item-count", fmt.Sprintf("%d", len(triggers)))
 		c.IndentedJSON(http.StatusOK, gin.H{"_rid": "", "Triggers": triggers, "_count": len(triggers)})
 		return
 	}

api/handlers/user_defined_functions.go

@@ -1,6 +1,7 @@
 package handlers
 
 import (
+	"fmt"
 	"net/http"
 
 	"github.com/gin-gonic/gin"
@@ -15,6 +16,7 @@ func GetAllUserDefinedFunctions(c *gin.Context) {
 	udfs, status := repositories.GetAllUserDefinedFunctions(databaseId, collectionId)
 
 	if status == repositorymodels.StatusOk {
+		c.Header("x-ms-item-count", fmt.Sprintf("%d", len(udfs)))
 		c.IndentedJSON(http.StatusOK, gin.H{"_rid": "", "UserDefinedFunctions": udfs, "_count": len(udfs)})
 		return
 	}

internal/authentication/authentication.go

@@ -10,6 +10,11 @@ import (
 
 // https://learn.microsoft.com/en-us/rest/api/cosmos-db/access-control-on-cosmosdb-resources
 func GenerateSignature(verb string, resourceType string, resourceId string, date string, masterKey string) string {
+	isNameBased := resourceId != "" && ((len(resourceId) > 4 && resourceId[3] == '/') || strings.HasPrefix(strings.ToLower(resourceId), "interopusers"))
+	if !isNameBased {
+		resourceId = strings.ToLower(resourceId)
+	}
+
 	payload := fmt.Sprintf(
 		"%s\n%s\n%s\n%s\n%s\n",
 		strings.ToLower(verb),

internal/authentication/authentication_test.go

@@ -27,4 +27,9 @@ func Test_GenerateSignature(t *testing.T) {
 		signature := authentication.GenerateSignature("DELETE", "dbs", "dbs/Test Database", testDate, config.DefaultAccountKey)
 		assert.Equal(t, "LcuXXg0TcXxZG0kUCj9tZIWRy2yCzim3oiqGiHpRqGs=", signature)
 	})
+
+	t.Run("Should generate PKRANGES signature", func(t *testing.T) {
+		signature := authentication.GenerateSignature("GET", "pkranges", "m4d+xG08uVM=", testDate, config.DefaultAccountKey)
+		assert.Equal(t, "6S5ceZsl2EXWB3Jo5bJcK7zv8NxXnsxWPWD9TH3nNMo=", signature)
+	})
 }

internal/repositories/documents.go

@@ -66,7 +66,8 @@ func CreateDocument(databaseId string, collectionId string, document map[string]
 	var database repositorymodels.Database
 	var collection repositorymodels.Collection
 	if documentId, ok = document["id"].(string); !ok || documentId == "" {
-		return repositorymodels.Document{}, repositorymodels.BadRequest
+		documentId = fmt.Sprint(uuid.New())
+		document["id"] = documentId
 	}
 
 	if database, ok = storeState.Databases[databaseId]; !ok {

internal/resourceid/resourceid.go

@@ -2,6 +2,7 @@ package resourceid
 
 import (
 	"encoding/base64"
+	"math/rand"
 
 	"github.com/google/uuid"
 )
@@ -10,6 +11,12 @@ func New() string {
 	id := uuid.New().ID()
 	idBytes := uintToBytes(id)
 
+	// first byte should be bigger than 0x80 for collection ids
+	// clients classify this id as "user" otherwise
+	if (idBytes[0] & 0x80) <= 0 {
+		idBytes[0] = byte(rand.Intn(0x80) + 0x80)
+	}
+
 	return base64.StdEncoding.EncodeToString(idBytes)
 }