From efcf4ad13d1af9cef15c480e68d73f58ba6aeb4c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=A5=BA=E5=AD=90w?= Date: Sat, 16 Mar 2019 18:10:28 +0800 Subject: [PATCH] Use crypto/rand for token generation (#161) --- api/application.go | 6 +++--- api/application_test.go | 39 +++++++++++++++++++++++++++------------ api/client.go | 2 +- api/client_test.go | 11 +++++++---- api/plugin_test.go | 2 -- api/tokens.go | 11 +++++++++++ api/tokens_test.go | 14 ++++++++++++++ auth/token.go | 29 +++++++++++++++++++++++------ auth/token_test.go | 13 +++++++++++++ test/token.go | 16 ++++++++++++++++ test/token_test.go | 15 +++++++++++++++ 11 files changed, 130 insertions(+), 28 deletions(-) create mode 100644 api/tokens.go create mode 100644 api/tokens_test.go create mode 100644 test/token.go create mode 100644 test/token_test.go diff --git a/api/application.go b/api/application.go index 6a501a3..a236fc0 100644 --- a/api/application.go +++ b/api/application.go @@ -65,7 +65,7 @@ type ApplicationAPI struct { func (a *ApplicationAPI) CreateApplication(ctx *gin.Context) { app := model.Application{} if err := ctx.Bind(&app); err == nil { - app.Token = auth.GenerateNotExistingToken(auth.GenerateApplicationToken, a.applicationExists) + app.Token = auth.GenerateNotExistingToken(generateApplicationToken, a.applicationExists) app.UserID = auth.GetUserID(ctx) app.Internal = false a.DB.CreateApplication(&app) @@ -284,9 +284,9 @@ func (a *ApplicationAPI) UploadApplicationImage(ctx *gin.Context) { ext := filepath.Ext(file.Filename) - name := auth.GenerateImageName() + name := generateImageName() for exist(a.ImageDir + name + ext) { - name = auth.GenerateImageName() + name = generateImageName() } err = ctx.SaveUploadedFile(file, a.ImageDir+name+ext) diff --git a/api/application_test.go b/api/application_test.go index 766619a..e41c033 100644 --- a/api/application_test.go +++ b/api/application_test.go @@ -5,7 +5,6 @@ import ( "errors" "io" "io/ioutil" - "math/rand" "mime/multipart" "net/http/httptest" "os" @@ -22,8 +21,8 @@ import ( ) var ( - firstApplicationToken = "APorrUa5b1IIK3y" - secondApplicationToken = "AKo_Pp6ww_9vZal" + firstApplicationToken = "Aaaaaaaaaaaaaaa" + secondApplicationToken = "Abbbbbbbbbbbbbb" ) func TestApplicationSuite(t *testing.T) { @@ -38,9 +37,15 @@ type ApplicationSuite struct { recorder *httptest.ResponseRecorder } +var originalGenerateApplicationToken func() string +var originalGenerateImageName func() string + func (s *ApplicationSuite) BeforeTest(suiteName, testName string) { + originalGenerateApplicationToken = generateApplicationToken + originalGenerateImageName = generateImageName + generateApplicationToken = test.Tokens(firstApplicationToken, secondApplicationToken) + generateImageName = test.Tokens(firstApplicationToken[1:], secondApplicationToken[1:]) mode.Set(mode.TestDev) - rand.Seed(50) s.recorder = httptest.NewRecorder() s.db = testdb.NewDB(s.T()) s.ctx, _ = gin.CreateTestContext(s.recorder) @@ -49,6 +54,8 @@ func (s *ApplicationSuite) BeforeTest(suiteName, testName string) { } func (s *ApplicationSuite) AfterTest(suiteName, testName string) { + generateApplicationToken = originalGenerateApplicationToken + generateImageName = originalGenerateImageName s.db.Close() } @@ -268,19 +275,24 @@ func (s *ApplicationSuite) Test_UploadAppImage_WithImageFile_expectSuccess() { s.a.UploadApplicationImage(s.ctx) + imgName := s.db.GetApplicationByID(1).Image + assert.Equal(s.T(), 200, s.recorder.Code) - _, err = os.Stat("PorrUa5b1IIK3yKo_Pp6ww_9v.png") + _, err = os.Stat(imgName) assert.Nil(s.T(), err) s.a.DeleteApplication(s.ctx) - _, err = os.Stat("PorrUa5b1IIK3yKo_Pp6ww_9v.png") + _, err = os.Stat(imgName) assert.True(s.T(), os.IsNotExist(err)) } func (s *ApplicationSuite) Test_UploadAppImage_WithImageFile_DeleteExstingImageAndGenerateNewName() { + existingImageName := "2lHMAel6BDHLL-HrwphcviX-l.png" + firstGeneratedImageName := firstApplicationToken[1:] + ".png" + secondGeneratedImageName := secondApplicationToken[1:] + ".png" s.db.User(5) - s.db.CreateApplication(&model.Application{UserID: 5, ID: 1, Image: "PorrUa5b1IIK3yKo_Pp6ww_9v.png"}) + s.db.CreateApplication(&model.Application{UserID: 5, ID: 1, Image: existingImageName}) cType, buffer, err := upload(map[string]*os.File{"file": mustOpen("../test/assets/image.png")}) assert.Nil(s.T(), err) @@ -288,17 +300,20 @@ func (s *ApplicationSuite) Test_UploadAppImage_WithImageFile_DeleteExstingImageA s.ctx.Request.Header.Set("Content-Type", cType) test.WithUser(s.ctx, 5) s.ctx.Params = gin.Params{{Key: "id", Value: "1"}} - fakeImage(s.T(), "PorrUa5b1IIK3yKo_Pp6ww_9v.png") + fakeImage(s.T(), existingImageName) + fakeImage(s.T(), firstGeneratedImageName) s.a.UploadApplicationImage(s.ctx) assert.Equal(s.T(), 200, s.recorder.Code) - _, err = os.Stat("PorrUa5b1IIK3yKo_Pp6ww_9v.png") + _, err = os.Stat(existingImageName) assert.True(s.T(), os.IsNotExist(err)) - _, err = os.Stat("Zal6-ySIuL-T3EMLCcFtityHn.png") + + _, err = os.Stat(secondGeneratedImageName) assert.Nil(s.T(), err) - assert.Nil(s.T(), os.Remove("Zal6-ySIuL-T3EMLCcFtityHn.png")) + assert.Nil(s.T(), os.Remove(secondGeneratedImageName)) + assert.Nil(s.T(), os.Remove(firstGeneratedImageName)) } func (s *ApplicationSuite) Test_UploadAppImage_WithImageFile_DeleteExistingImage() { @@ -320,7 +335,7 @@ func (s *ApplicationSuite) Test_UploadAppImage_WithImageFile_DeleteExistingImage _, err = os.Stat("existing.png") assert.True(s.T(), os.IsNotExist(err)) - os.Remove("PorrUa5b1IIK3yKo_Pp6ww_9v.png") + os.Remove(firstApplicationToken[1:] + ".png") } func (s *ApplicationSuite) Test_UploadAppImage_WithTextFile_expectBadRequest() { diff --git a/api/client.go b/api/client.go index 3f6fb7c..5a299f2 100644 --- a/api/client.go +++ b/api/client.go @@ -60,7 +60,7 @@ type ClientAPI struct { func (a *ClientAPI) CreateClient(ctx *gin.Context) { client := model.Client{} if err := ctx.Bind(&client); err == nil { - client.Token = auth.GenerateNotExistingToken(auth.GenerateClientToken, a.clientExists) + client.Token = auth.GenerateNotExistingToken(generateClientToken, a.clientExists) client.UserID = auth.GetUserID(ctx) a.DB.CreateClient(&client) ctx.JSON(200, client) diff --git a/api/client_test.go b/api/client_test.go index 9640121..32b1c8e 100644 --- a/api/client_test.go +++ b/api/client_test.go @@ -1,7 +1,6 @@ package api import ( - "math/rand" "net/http/httptest" "net/url" "strings" @@ -17,8 +16,8 @@ import ( ) var ( - firstClientToken = "CPorrUa5b1IIK3y" - secondClientToken = "CKo_Pp6ww_9vZal" + firstClientToken = "Caaaaaaaaaaaaaa" + secondClientToken = "Cbbbbbbbbbbbbbb" ) func TestClientSuite(t *testing.T) { @@ -34,9 +33,12 @@ type ClientSuite struct { notified bool } +var originalGenerateClientToken func() string + func (s *ClientSuite) BeforeTest(suiteName, testName string) { + originalGenerateClientToken = generateClientToken + generateClientToken = test.Tokens(firstClientToken, secondClientToken) mode.Set(mode.TestDev) - rand.Seed(50) s.recorder = httptest.NewRecorder() s.db = testdb.NewDB(s.T()) s.ctx, _ = gin.CreateTestContext(s.recorder) @@ -50,6 +52,7 @@ func (s *ClientSuite) notify(uint, string) { } func (s *ClientSuite) AfterTest(suiteName, testName string) { + generateClientToken = originalGenerateClientToken s.db.Close() } diff --git a/api/plugin_test.go b/api/plugin_test.go index 36eff0a..79f7114 100644 --- a/api/plugin_test.go +++ b/api/plugin_test.go @@ -5,7 +5,6 @@ import ( "encoding/json" "errors" "fmt" - "math/rand" "net/http/httptest" "testing" @@ -41,7 +40,6 @@ type PluginSuite struct { func (s *PluginSuite) BeforeTest(suiteName, testName string) { mode.Set(mode.TestDev) - rand.Seed(50) s.db = testdb.NewDB(s.T()) s.resetRecorder() manager, err := plugin.NewManager(s.db, "", nil, s) diff --git a/api/tokens.go b/api/tokens.go new file mode 100644 index 0000000..11591bc --- /dev/null +++ b/api/tokens.go @@ -0,0 +1,11 @@ +package api + +import ( + "github.com/gotify/server/auth" +) + +var generateApplicationToken = auth.GenerateApplicationToken + +var generateClientToken = auth.GenerateClientToken + +var generateImageName = auth.GenerateImageName diff --git a/api/tokens_test.go b/api/tokens_test.go new file mode 100644 index 0000000..66289e7 --- /dev/null +++ b/api/tokens_test.go @@ -0,0 +1,14 @@ +package api + +import ( + "regexp" + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestTokenGeneration(t *testing.T) { + assert.Regexp(t, regexp.MustCompile("^C(.+)$"), generateClientToken()) + assert.Regexp(t, regexp.MustCompile("^A(.+)$"), generateApplicationToken()) + assert.Regexp(t, regexp.MustCompile("^(.+)$"), generateImageName()) +} diff --git a/auth/token.go b/auth/token.go index 16cd3ff..f91185a 100644 --- a/auth/token.go +++ b/auth/token.go @@ -1,17 +1,29 @@ package auth import ( - "math/rand" + "crypto/rand" + "math/big" ) var ( - tokenCharacters = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789.-_") + tokenCharacters = []byte("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789.-_") randomTokenLength = 14 applicationPrefix = "A" clientPrefix = "C" pluginPrefix = "P" + + randReader = rand.Reader ) +func randIntn(n int) int { + max := big.NewInt(int64(n)) + res, err := rand.Int(randReader, max) + if err != nil { + panic("random source is not available") + } + return int(res.Int64()) +} + // GenerateNotExistingToken receives a token generation func and a func to check whether the token exists, returns a unique token. func GenerateNotExistingToken(generateToken func() string, tokenExists func(token string) bool) string { for { @@ -47,9 +59,14 @@ func generateRandomToken(prefix string) string { } func generateRandomString(length int) string { - b := make([]rune, length) - for i := range b { - b[i] = tokenCharacters[rand.Intn(len(tokenCharacters))] + res := make([]byte, length) + for i := range res { + index := randIntn(len(tokenCharacters)) + res[i] = tokenCharacters[index] } - return string(b) + return string(res) +} + +func init() { + randIntn(2) } diff --git a/auth/token_test.go b/auth/token_test.go index 99b157b..85384c8 100644 --- a/auth/token_test.go +++ b/auth/token_test.go @@ -1,10 +1,13 @@ package auth import ( + "crypto/rand" "fmt" "strings" "testing" + "github.com/gotify/server/test" + "github.com/stretchr/testify/assert" ) @@ -30,3 +33,13 @@ func TestGenerateNotExistingToken(t *testing.T) { }) assert.Equal(t, "0", token) } + +func TestBadCryptoReaderPanics(t *testing.T) { + assert.Panics(t, func() { + randReader = test.UnreadableReader() + defer func() { + randReader = rand.Reader + }() + randIntn(2) + }) +} diff --git a/test/token.go b/test/token.go new file mode 100644 index 0000000..64370f4 --- /dev/null +++ b/test/token.go @@ -0,0 +1,16 @@ +package test + +import "sync" + +// Tokens returns a token generation function with takes a series of tokens and output them in order. +func Tokens(tokens ...string) func() string { + var i int + lock := sync.Mutex{} + return func() string { + lock.Lock() + defer lock.Unlock() + res := tokens[i%len(tokens)] + i++ + return res + } +} diff --git a/test/token_test.go b/test/token_test.go new file mode 100644 index 0000000..000be3d --- /dev/null +++ b/test/token_test.go @@ -0,0 +1,15 @@ +package test + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestTokenGeneration(t *testing.T) { + mockTokenFunc := Tokens("a", "b", "c") + + for _, expected := range []string{"a", "b", "c", "a", "b", "c"} { + assert.Equal(t, expected, mockTokenFunc()) + } +}