Skip to content

Commit efcf4ad

Browse files
饺子wjmattheis
authored andcommitted
Use crypto/rand for token generation (#161)
1 parent 79b3a0c commit efcf4ad

File tree

11 files changed

+130
-28
lines changed

11 files changed

+130
-28
lines changed

api/application.go

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -65,7 +65,7 @@ type ApplicationAPI struct {
6565
func (a *ApplicationAPI) CreateApplication(ctx *gin.Context) {
6666
app := model.Application{}
6767
if err := ctx.Bind(&app); err == nil {
68-
app.Token = auth.GenerateNotExistingToken(auth.GenerateApplicationToken, a.applicationExists)
68+
app.Token = auth.GenerateNotExistingToken(generateApplicationToken, a.applicationExists)
6969
app.UserID = auth.GetUserID(ctx)
7070
app.Internal = false
7171
a.DB.CreateApplication(&app)
@@ -284,9 +284,9 @@ func (a *ApplicationAPI) UploadApplicationImage(ctx *gin.Context) {
284284

285285
ext := filepath.Ext(file.Filename)
286286

287-
name := auth.GenerateImageName()
287+
name := generateImageName()
288288
for exist(a.ImageDir + name + ext) {
289-
name = auth.GenerateImageName()
289+
name = generateImageName()
290290
}
291291

292292
err = ctx.SaveUploadedFile(file, a.ImageDir+name+ext)

api/application_test.go

Lines changed: 27 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@ import (
55
"errors"
66
"io"
77
"io/ioutil"
8-
"math/rand"
98
"mime/multipart"
109
"net/http/httptest"
1110
"os"
@@ -22,8 +21,8 @@ import (
2221
)
2322

2423
var (
25-
firstApplicationToken = "APorrUa5b1IIK3y"
26-
secondApplicationToken = "AKo_Pp6ww_9vZal"
24+
firstApplicationToken = "Aaaaaaaaaaaaaaa"
25+
secondApplicationToken = "Abbbbbbbbbbbbbb"
2726
)
2827

2928
func TestApplicationSuite(t *testing.T) {
@@ -38,9 +37,15 @@ type ApplicationSuite struct {
3837
recorder *httptest.ResponseRecorder
3938
}
4039

40+
var originalGenerateApplicationToken func() string
41+
var originalGenerateImageName func() string
42+
4143
func (s *ApplicationSuite) BeforeTest(suiteName, testName string) {
44+
originalGenerateApplicationToken = generateApplicationToken
45+
originalGenerateImageName = generateImageName
46+
generateApplicationToken = test.Tokens(firstApplicationToken, secondApplicationToken)
47+
generateImageName = test.Tokens(firstApplicationToken[1:], secondApplicationToken[1:])
4248
mode.Set(mode.TestDev)
43-
rand.Seed(50)
4449
s.recorder = httptest.NewRecorder()
4550
s.db = testdb.NewDB(s.T())
4651
s.ctx, _ = gin.CreateTestContext(s.recorder)
@@ -49,6 +54,8 @@ func (s *ApplicationSuite) BeforeTest(suiteName, testName string) {
4954
}
5055

5156
func (s *ApplicationSuite) AfterTest(suiteName, testName string) {
57+
generateApplicationToken = originalGenerateApplicationToken
58+
generateImageName = originalGenerateImageName
5259
s.db.Close()
5360
}
5461

@@ -268,37 +275,45 @@ func (s *ApplicationSuite) Test_UploadAppImage_WithImageFile_expectSuccess() {
268275

269276
s.a.UploadApplicationImage(s.ctx)
270277

278+
imgName := s.db.GetApplicationByID(1).Image
279+
271280
assert.Equal(s.T(), 200, s.recorder.Code)
272-
_, err = os.Stat("PorrUa5b1IIK3yKo_Pp6ww_9v.png")
281+
_, err = os.Stat(imgName)
273282
assert.Nil(s.T(), err)
274283

275284
s.a.DeleteApplication(s.ctx)
276285

277-
_, err = os.Stat("PorrUa5b1IIK3yKo_Pp6ww_9v.png")
286+
_, err = os.Stat(imgName)
278287
assert.True(s.T(), os.IsNotExist(err))
279288
}
280289

281290
func (s *ApplicationSuite) Test_UploadAppImage_WithImageFile_DeleteExstingImageAndGenerateNewName() {
291+
existingImageName := "2lHMAel6BDHLL-HrwphcviX-l.png"
292+
firstGeneratedImageName := firstApplicationToken[1:] + ".png"
293+
secondGeneratedImageName := secondApplicationToken[1:] + ".png"
282294
s.db.User(5)
283-
s.db.CreateApplication(&model.Application{UserID: 5, ID: 1, Image: "PorrUa5b1IIK3yKo_Pp6ww_9v.png"})
295+
s.db.CreateApplication(&model.Application{UserID: 5, ID: 1, Image: existingImageName})
284296

285297
cType, buffer, err := upload(map[string]*os.File{"file": mustOpen("../test/assets/image.png")})
286298
assert.Nil(s.T(), err)
287299
s.ctx.Request = httptest.NewRequest("POST", "/irrelevant", &buffer)
288300
s.ctx.Request.Header.Set("Content-Type", cType)
289301
test.WithUser(s.ctx, 5)
290302
s.ctx.Params = gin.Params{{Key: "id", Value: "1"}}
291-
fakeImage(s.T(), "PorrUa5b1IIK3yKo_Pp6ww_9v.png")
303+
fakeImage(s.T(), existingImageName)
304+
fakeImage(s.T(), firstGeneratedImageName)
292305

293306
s.a.UploadApplicationImage(s.ctx)
294307

295308
assert.Equal(s.T(), 200, s.recorder.Code)
296309

297-
_, err = os.Stat("PorrUa5b1IIK3yKo_Pp6ww_9v.png")
310+
_, err = os.Stat(existingImageName)
298311
assert.True(s.T(), os.IsNotExist(err))
299-
_, err = os.Stat("Zal6-ySIuL-T3EMLCcFtityHn.png")
312+
313+
_, err = os.Stat(secondGeneratedImageName)
300314
assert.Nil(s.T(), err)
301-
assert.Nil(s.T(), os.Remove("Zal6-ySIuL-T3EMLCcFtityHn.png"))
315+
assert.Nil(s.T(), os.Remove(secondGeneratedImageName))
316+
assert.Nil(s.T(), os.Remove(firstGeneratedImageName))
302317
}
303318

304319
func (s *ApplicationSuite) Test_UploadAppImage_WithImageFile_DeleteExistingImage() {
@@ -320,7 +335,7 @@ func (s *ApplicationSuite) Test_UploadAppImage_WithImageFile_DeleteExistingImage
320335
_, err = os.Stat("existing.png")
321336
assert.True(s.T(), os.IsNotExist(err))
322337

323-
os.Remove("PorrUa5b1IIK3yKo_Pp6ww_9v.png")
338+
os.Remove(firstApplicationToken[1:] + ".png")
324339
}
325340

326341
func (s *ApplicationSuite) Test_UploadAppImage_WithTextFile_expectBadRequest() {

api/client.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,7 @@ type ClientAPI struct {
6060
func (a *ClientAPI) CreateClient(ctx *gin.Context) {
6161
client := model.Client{}
6262
if err := ctx.Bind(&client); err == nil {
63-
client.Token = auth.GenerateNotExistingToken(auth.GenerateClientToken, a.clientExists)
63+
client.Token = auth.GenerateNotExistingToken(generateClientToken, a.clientExists)
6464
client.UserID = auth.GetUserID(ctx)
6565
a.DB.CreateClient(&client)
6666
ctx.JSON(200, client)

api/client_test.go

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
package api
22

33
import (
4-
"math/rand"
54
"net/http/httptest"
65
"net/url"
76
"strings"
@@ -17,8 +16,8 @@ import (
1716
)
1817

1918
var (
20-
firstClientToken = "CPorrUa5b1IIK3y"
21-
secondClientToken = "CKo_Pp6ww_9vZal"
19+
firstClientToken = "Caaaaaaaaaaaaaa"
20+
secondClientToken = "Cbbbbbbbbbbbbbb"
2221
)
2322

2423
func TestClientSuite(t *testing.T) {
@@ -34,9 +33,12 @@ type ClientSuite struct {
3433
notified bool
3534
}
3635

36+
var originalGenerateClientToken func() string
37+
3738
func (s *ClientSuite) BeforeTest(suiteName, testName string) {
39+
originalGenerateClientToken = generateClientToken
40+
generateClientToken = test.Tokens(firstClientToken, secondClientToken)
3841
mode.Set(mode.TestDev)
39-
rand.Seed(50)
4042
s.recorder = httptest.NewRecorder()
4143
s.db = testdb.NewDB(s.T())
4244
s.ctx, _ = gin.CreateTestContext(s.recorder)
@@ -50,6 +52,7 @@ func (s *ClientSuite) notify(uint, string) {
5052
}
5153

5254
func (s *ClientSuite) AfterTest(suiteName, testName string) {
55+
generateClientToken = originalGenerateClientToken
5356
s.db.Close()
5457
}
5558

api/plugin_test.go

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@ import (
55
"encoding/json"
66
"errors"
77
"fmt"
8-
"math/rand"
98
"net/http/httptest"
109
"testing"
1110

@@ -41,7 +40,6 @@ type PluginSuite struct {
4140

4241
func (s *PluginSuite) BeforeTest(suiteName, testName string) {
4342
mode.Set(mode.TestDev)
44-
rand.Seed(50)
4543
s.db = testdb.NewDB(s.T())
4644
s.resetRecorder()
4745
manager, err := plugin.NewManager(s.db, "", nil, s)

api/tokens.go

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
package api
2+
3+
import (
4+
"github.com/gotify/server/auth"
5+
)
6+
7+
var generateApplicationToken = auth.GenerateApplicationToken
8+
9+
var generateClientToken = auth.GenerateClientToken
10+
11+
var generateImageName = auth.GenerateImageName

api/tokens_test.go

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
package api
2+
3+
import (
4+
"regexp"
5+
"testing"
6+
7+
"github.com/stretchr/testify/assert"
8+
)
9+
10+
func TestTokenGeneration(t *testing.T) {
11+
assert.Regexp(t, regexp.MustCompile("^C(.+)$"), generateClientToken())
12+
assert.Regexp(t, regexp.MustCompile("^A(.+)$"), generateApplicationToken())
13+
assert.Regexp(t, regexp.MustCompile("^(.+)$"), generateImageName())
14+
}

auth/token.go

Lines changed: 23 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,29 @@
11
package auth
22

33
import (
4-
"math/rand"
4+
"crypto/rand"
5+
"math/big"
56
)
67

78
var (
8-
tokenCharacters = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789.-_")
9+
tokenCharacters = []byte("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789.-_")
910
randomTokenLength = 14
1011
applicationPrefix = "A"
1112
clientPrefix = "C"
1213
pluginPrefix = "P"
14+
15+
randReader = rand.Reader
1316
)
1417

18+
func randIntn(n int) int {
19+
max := big.NewInt(int64(n))
20+
res, err := rand.Int(randReader, max)
21+
if err != nil {
22+
panic("random source is not available")
23+
}
24+
return int(res.Int64())
25+
}
26+
1527
// GenerateNotExistingToken receives a token generation func and a func to check whether the token exists, returns a unique token.
1628
func GenerateNotExistingToken(generateToken func() string, tokenExists func(token string) bool) string {
1729
for {
@@ -47,9 +59,14 @@ func generateRandomToken(prefix string) string {
4759
}
4860

4961
func generateRandomString(length int) string {
50-
b := make([]rune, length)
51-
for i := range b {
52-
b[i] = tokenCharacters[rand.Intn(len(tokenCharacters))]
62+
res := make([]byte, length)
63+
for i := range res {
64+
index := randIntn(len(tokenCharacters))
65+
res[i] = tokenCharacters[index]
5366
}
54-
return string(b)
67+
return string(res)
68+
}
69+
70+
func init() {
71+
randIntn(2)
5572
}

auth/token_test.go

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,13 @@
11
package auth
22

33
import (
4+
"crypto/rand"
45
"fmt"
56
"strings"
67
"testing"
78

9+
"github.com/gotify/server/test"
10+
811
"github.com/stretchr/testify/assert"
912
)
1013

@@ -30,3 +33,13 @@ func TestGenerateNotExistingToken(t *testing.T) {
3033
})
3134
assert.Equal(t, "0", token)
3235
}
36+
37+
func TestBadCryptoReaderPanics(t *testing.T) {
38+
assert.Panics(t, func() {
39+
randReader = test.UnreadableReader()
40+
defer func() {
41+
randReader = rand.Reader
42+
}()
43+
randIntn(2)
44+
})
45+
}

test/token.go

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
package test
2+
3+
import "sync"
4+
5+
// Tokens returns a token generation function with takes a series of tokens and output them in order.
6+
func Tokens(tokens ...string) func() string {
7+
var i int
8+
lock := sync.Mutex{}
9+
return func() string {
10+
lock.Lock()
11+
defer lock.Unlock()
12+
res := tokens[i%len(tokens)]
13+
i++
14+
return res
15+
}
16+
}

0 commit comments

Comments
 (0)