use tokenMarshal/unmarshal in page encoding

Signed-off-by: liangchenye <liangchenye@huawei.com>

api/v1/models.go

@@ -19,7 +19,6 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
-	"strconv"
 	"time"
 
 	"github.com/coreos/clair/database"
@@ -216,7 +215,8 @@ func NotificationFromDatabaseModel(dbNotification database.VulnerabilityNotifica
 
 	var nextPageStr string
 	if nextPage != database.NoVulnerabilityNotificationPage {
-		nextPageStr = pageNumberToToken(nextPage, key)
+		nextPageBytes, _ := tokenMarshal(nextPage, key)
+		nextPageStr = string(nextPageBytes)
 	}
 
 	var created, notified, deleted string
@@ -292,51 +292,23 @@ type FeatureEnvelope struct {
 	Error    *Error     `json:"Error,omitempty"`
 }
 
-func tokenToPageNumber(token, key string) (database.VulnerabilityNotificationPageNumber, error) {
+func tokenUnmarshal(token string, key string, v interface{}) error {
 	k, _ := fernet.DecodeKey(key)
 	msg := fernet.VerifyAndDecrypt([]byte(token), time.Hour, []*fernet.Key{k})
 	if msg == nil {
-		return database.VulnerabilityNotificationPageNumber{}, errors.New("invalid or expired pagination token")
+		return errors.New("invalid or expired pagination token")
 	}
 
-	page := database.VulnerabilityNotificationPageNumber{}
-	err := json.NewDecoder(bytes.NewBuffer(msg)).Decode(&page)
-	return page, err
+	return json.NewDecoder(bytes.NewBuffer(msg)).Decode(&v)
 }
 
-func pageNumberToToken(page database.VulnerabilityNotificationPageNumber, key string) string {
+func tokenMarshal(v interface{}, key string) ([]byte, error) {
 	var buf bytes.Buffer
-	err := json.NewEncoder(&buf).Encode(page)
+	err := json.NewEncoder(&buf).Encode(v)
 	if err != nil {
-		log.Fatal("failed to encode VulnerabilityNotificationPageNumber")
+		return nil, err
 	}
 
 	k, _ := fernet.DecodeKey(key)
-	tokenBytes, err := fernet.EncryptAndSign(buf.Bytes(), k)
-	if err != nil {
-		log.Fatal("failed to encrypt VulnerabilityNotificationpageNumber")
-	}
-
-	return string(tokenBytes)
-}
-
-func tokenToNumber(token, key string) (int, error) {
-	k, _ := fernet.DecodeKey(key)
-	msg := fernet.VerifyAndDecrypt([]byte(token), time.Hour, []*fernet.Key{k})
-	if msg == nil {
-		return -1, errors.New("invalid or expired pagination token")
-	}
-
-	page, err := strconv.Atoi(string(msg))
-	return page, err
-}
-
-func numberToToken(page int, key string) string {
-	k, _ := fernet.DecodeKey(key)
-	tokenBytes, err := fernet.EncryptAndSign([]byte(strconv.Itoa(page)), k)
-	if err != nil {
-		log.Fatal("failed to encrypt number")
-	}
-
-	return string(tokenBytes)
+	return fernet.EncryptAndSign(buf.Bytes(), k)
 }

api/v1/routes.go

@@ -197,7 +197,7 @@ func getVulnerabilities(w http.ResponseWriter, r *http.Request, p httprouter.Par
 	if err != nil {
 		writeResponse(w, r, http.StatusBadRequest, VulnerabilityEnvelope{Error: &Error{"invalid limit format: " + err.Error()}})
 		return getVulnerabilitiesRoute, http.StatusBadRequest
-	} else if limit <= 0 {
+	} else if limit < 0 {
 		writeResponse(w, r, http.StatusBadRequest, VulnerabilityEnvelope{Error: &Error{"limit value should not be less than zero"}})
 		return getVulnerabilitiesRoute, http.StatusBadRequest
 	}
@@ -205,7 +205,7 @@ func getVulnerabilities(w http.ResponseWriter, r *http.Request, p httprouter.Par
 	page := 0
 	pageStrs, pageExists := query["page"]
 	if pageExists {
-		page, err = tokenToNumber(pageStrs[0], ctx.Config.PaginationKey)
+		err = tokenUnmarshal(pageStrs[0], ctx.Config.PaginationKey, &page)
 		if err != nil {
 			writeResponse(w, r, http.StatusBadRequest, VulnerabilityEnvelope{Error: &Error{"invalid page format: " + err.Error()}})
 			return getNotificationRoute, http.StatusBadRequest
@@ -226,7 +226,12 @@ func getVulnerabilities(w http.ResponseWriter, r *http.Request, p httprouter.Par
 
 	var nextPageStr string
 	if nextPage != -1 {
-		nextPageStr = numberToToken(nextPage, ctx.Config.PaginationKey)
+		nextPageBytes, err := tokenMarshal(nextPage, ctx.Config.PaginationKey)
+		if err != nil {
+			writeResponse(w, r, http.StatusBadRequest, VulnerabilityEnvelope{Error: &Error{"failed to marshal token: " + err.Error()}})
+			return getNotificationRoute, http.StatusBadRequest
+		}
+		nextPageStr = string(nextPageBytes)
 	}
 
 	writeResponse(w, r, http.StatusOK, VulnerabilityEnvelope{Vulnerabilities: &vulns, NextPage: &nextPageStr})
@@ -434,14 +439,19 @@ func getNotification(w http.ResponseWriter, r *http.Request, p httprouter.Params
 	page := database.VulnerabilityNotificationFirstPage
 	pageStrs, pageExists := query["page"]
 	if pageExists {
-		page, err = tokenToPageNumber(pageStrs[0], ctx.Config.PaginationKey)
+		err := tokenUnmarshal(pageStrs[0], ctx.Config.PaginationKey, &page)
 		if err != nil {
 			writeResponse(w, r, http.StatusBadRequest, NotificationEnvelope{Error: &Error{"invalid page format: " + err.Error()}})
 			return getNotificationRoute, http.StatusBadRequest
 		}
 		pageToken = pageStrs[0]
 	} else {
-		pageToken = pageNumberToToken(page, ctx.Config.PaginationKey)
+		pageTokenBytes, err := tokenMarshal(page, ctx.Config.PaginationKey)
+		if err != nil {
+			writeResponse(w, r, http.StatusBadRequest, NotificationEnvelope{Error: &Error{"failed to marshal token: " + err.Error()}})
+			return getNotificationRoute, http.StatusBadRequest
+		}
+		pageToken = string(pageTokenBytes)
 	}
 
 	dbNotification, nextPage, err := ctx.Store.GetNotification(p.ByName("notificationName"), limit, page)