Golang rest.AppSimple函数代码示例

// TestGJRMiddlewareWithRaceCondition test GRJ middleware under race condition.
func TestGJRMiddlewareWithRaceCondition(t *testing.T) {
	runtime.GOMAXPROCS(4)

	api := rest.NewApi()

	api.Use(NewGJRMiddleware(newRedisLimiter("5-M", "limitertests:gjrrace")))

	api.SetApp(rest.AppSimple(func(w rest.ResponseWriter, r *rest.Request) {
		w.WriteJson(map[string]string{"message": "ok"})
	}))

	handler := api.MakeHandler()
	req := test.MakeSimpleRequest("GET", "http://localhost/", nil)
	req.RemoteAddr = fmt.Sprintf("178.1.2.%d:180", Random(1, 90))

	nbRequests := 100
	successCount := 0

	var wg sync.WaitGroup
	wg.Add(nbRequests)

	for i := 1; i <= nbRequests; i++ {
		go func() {
			recorded := test.RunRequest(t, handler, req)
			if recorded.Recorder.Code == 200 {
				successCount++
			}
			wg.Done()
		}()
	}

	wg.Wait()

	assert.Equal(t, 5, successCount)
}

func main() {
	api := rest.NewApi()
	api.Use(rest.DefaultDevStack...)
	api.SetApp(rest.AppSimple(func(w rest.ResponseWriter, r *rest.Request) {

		conf := os.Getenv("CONFIG")
		app := os.Getenv("APP")

		w.Header().Set("Content-Type", "text/plain")
		cmd := exec.Command("transporter", "run", "--config", conf, app)
		outs, err := cmd.CombinedOutput()

		out_cmds := fmt.Sprintf("==> Executing: %s\n", strings.Join(cmd.Args, " "))
		w.(http.ResponseWriter).Write([]byte(out_cmds))

		if err != nil {
			out_errs := fmt.Sprintf("==> Error: %s\n", err.Error())
			w.(http.ResponseWriter).Write([]byte(out_errs))
		} else {
			out_outs := fmt.Sprintf("==> Output: %s\n", string(outs))
			w.(http.ResponseWriter).Write([]byte(out_outs))
		}

	}))
	log.Fatal(http.ListenAndServe(":8080", api.MakeHandler()))
}

// TestRate tests ratelimit.Rate methods.
func TestGJRMiddleware(t *testing.T) {
	api := rest.NewApi()

	api.Use(NewGJRMiddleware(newRedisLimiter("10-M", "limitertests:gjr")))

	var reset int64

	api.SetApp(rest.AppSimple(func(w rest.ResponseWriter, r *rest.Request) {
		reset = r.Env["ratelimit:reset"].(int64)
		w.WriteJson(map[string]string{"message": "ok"})
	}))

	handler := api.MakeHandler()
	req := test.MakeSimpleRequest("GET", "http://localhost/", nil)
	req.RemoteAddr = fmt.Sprintf("178.1.2.%d:120", Random(1, 90))

	i := 1
	for i < 20 {
		recorded := test.RunRequest(t, handler, req)
		assert.True(t, math.Ceil(time.Since(time.Unix(reset, 0)).Seconds()) <= 60)
		if i <= 10 {
			recorded.BodyIs(`{"message":"ok"}`)
			recorded.HeaderIs("X-Ratelimit-Limit", "10")
			recorded.HeaderIs("X-Ratelimit-Remaining", fmt.Sprintf("%d", 10-i))
			recorded.CodeIs(200)
		} else {
			recorded.BodyIs(`{"Error":"Limit exceeded"}`)
			recorded.HeaderIs("X-Ratelimit-Limit", "10")
			recorded.HeaderIs("X-Ratelimit-Remaining", "0")
			recorded.CodeIs(429)
		}
		i++
	}
}

func main() {
	api := rest.NewApi()
	api.Use(rest.DefaultDevStack...)
	api.SetApp(rest.AppSimple(func(w rest.ResponseWriter, r *rest.Request) {
		w.WriteJson(map[string]string{"Body": "Hello World!"})
	}))
	log.Fatal(http.ListenAndServe(":8080", api.MakeHandler()))
}

func TestRefreshJWT(t *testing.T) {
	authMiddleware := &JWTMiddleware{
		Realm:      "test zone",
		Key:        key,
		Timeout:    time.Hour,
		MaxRefresh: time.Hour * 24,
		Authenticator: func(userId string, password string) bool {
			return userId == "admin" && password == "admin"
		},
		Authorizator: func(userId string, request *rest.Request) bool {
			return request.Method == "GET"
		},
	}

	loginApi := rest.NewApi()
	loginApi.SetApp(rest.AppSimple(authMiddleware.RefreshHandler))
	loginApi.Use(&rest.IfMiddleware{
		// Only authenticate non /refresh requests
		Condition: func(request *rest.Request) bool {
			return request.URL.Path != "/refresh"
		},
		IfTrue: authMiddleware,
	})
	api_router, _ := rest.MakeRouter(
		rest.Get("/refresh", authMiddleware.RefreshHandler),
	)
	loginApi.SetApp(api_router)

	// valid refresh (exp is >= Timeout and <= MaxRefresh, so it's refreshable)
	refreshableOutdatedToken := jwt.New(jwt.GetSigningMethod("HS256"))
	refreshableOutdatedToken.Claims["id"] = "admin"
	// we need to substract one hour and 1 sec to make token expired
	refreshableOutdatedToken.Claims["exp"] = time.Now().Unix() - 1
	refreshableOutdatedToken.Claims["orig_iat"] = time.Now().Add(-time.Hour).Unix() - 1
	tokenString, _ := refreshableOutdatedToken.SignedString(key)

	validRefreshOutdatedReq := test.MakeSimpleRequest("GET", "http://localhost/refresh", nil)
	validRefreshOutdatedReq.Header.Set("Authorization", "Bearer "+tokenString)

	recorded := test.RunRequest(t, loginApi.MakeHandler(), validRefreshOutdatedReq)
	recorded.CodeIs(200)
	recorded.ContentTypeIsJson()
}

func main() {
	// 4 reqs/hour
	rate, err := limiter.NewRateFromFormatted("4-H")
	if err != nil {
		panic(err)
	}

	// Create a Redis pool.
	pool := redis.NewPool(func() (redis.Conn, error) {
		c, err := redis.Dial("tcp", ":6379")
		if err != nil {
			return nil, err
		}
		return c, err
	}, 100)

	// Create a store with the pool.
	store, err := limiter.NewRedisStoreWithOptions(
		pool,
		limiter.StoreOptions{Prefix: "limiter_gjr_example", MaxRetry: 3})

	if err != nil {
		panic(err)
	}

	// Create API.
	api := rest.NewApi()
	api.Use(rest.DefaultDevStack...)

	// Add middleware with the limiter instance.
	api.Use(limiter.NewGJRMiddleware(limiter.NewLimiter(store, rate)))

	// Set stupid app.
	api.SetApp(rest.AppSimple(func(w rest.ResponseWriter, r *rest.Request) {
		w.WriteJson(map[string]string{"message": "ok"})
	}))

	// Run server!
	fmt.Println("Server is running on 7777...")
	log.Fatal(http.ListenAndServe(":7777", api.MakeHandler()))
}

func TestAuthJWTasymmetricKeys(t *testing.T) {
	privateKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		panic(err)
	}

	authMiddleware := &Middleware{
		Realm:            "test zone",
		SigningAlgorithm: "ES256",
		Key:              privateKey,
		VerifyKey:        &privateKey.PublicKey,
		Timeout:          time.Hour,
		MaxRefresh:       time.Hour * 24,
		PayloadFunc: func(userId string) map[string]interface{} {
			// Set custom claim, to be checked in Authorizator method
			return map[string]interface{}{"testkey": "testval", "exp": 0}
		},
		Authenticator: func(userId string, password string) error {
			// Not testing authentication, just authorization, so always return true
			return nil
		},
		Authorizator: func(userId string, request *rest.Request) bool {
			jwtClaims := ExtractClaims(request)

			// Check the actual claim, set in PayloadFunc
			return (jwtClaims["testkey"] == "testval")
		},
	}

	// Simple endpoint
	endpoint := func(w rest.ResponseWriter, r *rest.Request) {
		// Dummy endpoint, output doesn't really matter, we are checking
		// the code returned
		w.WriteJson(map[string]string{"Id": "123"})
	}

	// Setup simple app structure
	loginAPI := rest.NewApi()
	loginAPI.SetApp(rest.AppSimple(authMiddleware.LoginHandler))
	loginAPI.Use(&rest.IfMiddleware{
		// Only authenticate non /login requests
		Condition: func(request *rest.Request) bool {
			return request.URL.Path != "/login"
		},
		IfTrue: authMiddleware,
	})
	apiRouter, _ := rest.MakeRouter(
		rest.Post("/login", authMiddleware.LoginHandler),
		rest.Get("/", endpoint),
	)
	loginAPI.SetApp(apiRouter)

	// Authenticate
	loginCreds := map[string]string{"username": "admin", "password": "admin"}
	rightCredReq := test.MakeSimpleRequest("POST", "http://localhost/login", loginCreds)
	recorded := test.RunRequest(t, loginAPI.MakeHandler(), rightCredReq)
	recorded.CodeIs(200)
	recorded.ContentTypeIsJson()

	// Decode received token, to be sent with endpoint request
	nToken := DecoderToken{}
	test.DecodeJsonPayload(recorded.Recorder, &nToken)

	// Request endpoint, triggering Authorization.
	// If we get a 200 then the claims were available in Authorizator method
	req := test.MakeSimpleRequest("GET", "http://localhost/", nil)
	req.Header.Set("Authorization", "Bearer "+nToken.Token)
	recorded = test.RunRequest(t, loginAPI.MakeHandler(), req)
	recorded.CodeIs(200)
	recorded.ContentTypeIsJson()
}

func TestAuthJWT(t *testing.T) {

	// the middleware to test
	authMiddleware := &Middleware{
		Realm:      "test zone",
		Key:        key,
		Timeout:    time.Hour,
		MaxRefresh: time.Hour * 24,
		Authenticator: func(userId string, password string) error {
			if userId == "admin" && password == "admin" {
				return nil
			}
			return errors.New("Username and passwork should be admin")
		},
		Authorizator: func(userId string, request *rest.Request) bool {
			if request.Method == "GET" {
				return true
			}
			return false
		},
	}

	// api for testing failure
	apiFailure := rest.NewApi()
	apiFailure.Use(authMiddleware)
	apiFailure.SetApp(rest.AppSimple(func(w rest.ResponseWriter, r *rest.Request) {
		t.Error("Should never be executed")
	}))
	handler := apiFailure.MakeHandler()

	// simple request fails
	recorded := test.RunRequest(t, handler, test.MakeSimpleRequest("GET", "http://localhost/", nil))
	recorded.CodeIs(401)
	recorded.ContentTypeIsJson()

	// auth with right cred and wrong method fails
	wrongMethodReq := test.MakeSimpleRequest("POST", "http://localhost/", nil)
	wrongMethodReq.Header.Set("Authorization", "Bearer "+makeTokenString("admin", key))
	recorded = test.RunRequest(t, handler, wrongMethodReq)
	recorded.CodeIs(401)
	recorded.ContentTypeIsJson()

	// wrong Auth format - bearer lower case
	wrongAuthFormat := test.MakeSimpleRequest("GET", "http://localhost/", nil)
	wrongAuthFormat.Header.Set("Authorization", "bearer "+makeTokenString("admin", key))
	recorded = test.RunRequest(t, handler, wrongAuthFormat)
	recorded.CodeIs(401)
	recorded.ContentTypeIsJson()

	// wrong Auth format - no space after bearer
	wrongAuthFormat = test.MakeSimpleRequest("GET", "http://localhost/", nil)
	wrongAuthFormat.Header.Set("Authorization", "bearer"+makeTokenString("admin", key))
	recorded = test.RunRequest(t, handler, wrongAuthFormat)
	recorded.CodeIs(401)
	recorded.ContentTypeIsJson()

	// wrong Auth format - empty auth header
	wrongAuthFormat = test.MakeSimpleRequest("GET", "http://localhost/", nil)
	wrongAuthFormat.Header.Set("Authorization", "")
	recorded = test.RunRequest(t, handler, wrongAuthFormat)
	recorded.CodeIs(401)
	recorded.ContentTypeIsJson()

	// right credt, right method but wrong priv key
	wrongPrivKeyReq := test.MakeSimpleRequest("GET", "http://localhost/", nil)
	wrongPrivKeyReq.Header.Set("Authorization", "Bearer "+makeTokenString("admin", []byte("sekret key")))
	recorded = test.RunRequest(t, handler, wrongPrivKeyReq)
	recorded.CodeIs(401)
	recorded.ContentTypeIsJson()

	// right credt, right method, right priv key but timeout
	token := jwt.New(jwt.GetSigningMethod("HS256"))
	token.Claims["id"] = "admin"
	token.Claims["exp"] = 0
	tokenString, _ := token.SignedString(key)

	expiredTimestampReq := test.MakeSimpleRequest("GET", "http://localhost/", nil)
	expiredTimestampReq.Header.Set("Authorization", "Bearer "+tokenString)
	recorded = test.RunRequest(t, handler, expiredTimestampReq)
	recorded.CodeIs(401)
	recorded.ContentTypeIsJson()

	// right credt, right method, right priv, wrong signing method on request
	tokenBadSigning := jwt.New(jwt.GetSigningMethod("HS384"))
	tokenBadSigning.Claims["id"] = "admin"
	tokenBadSigning.Claims["exp"] = time.Now().Add(time.Hour * 72).Unix()
	tokenBadSigningString, _ := tokenBadSigning.SignedString(key)

	BadSigningReq := test.MakeSimpleRequest("GET", "http://localhost/", nil)
	BadSigningReq.Header.Set("Authorization", "Bearer "+tokenBadSigningString)
	recorded = test.RunRequest(t, handler, BadSigningReq)
	recorded.CodeIs(401)
	recorded.ContentTypeIsJson()

	// api for testing success
	apiSuccess := rest.NewApi()
	apiSuccess.Use(authMiddleware)
	apiSuccess.SetApp(rest.AppSimple(func(w rest.ResponseWriter, r *rest.Request) {
		if r.Env["REMOTE_USER"] == nil {
			t.Error("REMOTE_USER is nil")
//.........这里部分代码省略.........

func TestAuthJWTPayload(t *testing.T) {
	authMiddleware := &Middleware{
		Realm:            "test zone",
		SigningAlgorithm: "HS256",
		Key:              key,
		Timeout:          time.Hour,
		MaxRefresh:       time.Hour * 24,
		Authenticator: func(userId string, password string) error {
			if userId == "admin" && password == "admin" {
				return nil
			}
			return errors.New("Username and passwork should be admin")
		},
		PayloadFunc: func(userId string) map[string]interface{} {
			// tests normal value
			// tests overwriting of reserved jwt values should have no effect
			return map[string]interface{}{"testkey": "testval", "exp": 0}
		},
	}

	loginAPI := rest.NewApi()
	loginAPI.SetApp(rest.AppSimple(authMiddleware.LoginHandler))

	// correct payload
	loginCreds := map[string]string{"username": "admin", "password": "admin"}
	rightCredReq := test.MakeSimpleRequest("POST", "http://localhost/", loginCreds)
	recorded := test.RunRequest(t, loginAPI.MakeHandler(), rightCredReq)
	recorded.CodeIs(200)
	recorded.ContentTypeIsJson()

	nToken := DecoderToken{}
	test.DecodeJsonPayload(recorded.Recorder, &nToken)
	newToken, err := jwt.Parse(nToken.Token, func(token *jwt.Token) (interface{}, error) {
		return key, nil
	})

	if err != nil {
		t.Errorf("Received new token with wrong signature %s", err)
	}

	if newToken.Claims["testkey"].(string) != "testval" || newToken.Claims["exp"].(float64) == 0 {
		t.Errorf("Received new token without payload")
	}

	// correct payload after refresh
	refreshAPI := rest.NewApi()
	refreshAPI.Use(authMiddleware)
	refreshAPI.SetApp(rest.AppSimple(authMiddleware.RefreshHandler))

	refreshableToken := jwt.New(jwt.GetSigningMethod("HS256"))
	refreshableToken.Claims["id"] = "admin"
	refreshableToken.Claims["exp"] = time.Now().Add(time.Hour).Unix()
	refreshableToken.Claims["orig_iat"] = time.Now().Unix()
	refreshableToken.Claims["testkey"] = "testval"
	tokenString, _ := refreshableToken.SignedString(key)

	validRefreshReq := test.MakeSimpleRequest("GET", "http://localhost/", nil)
	validRefreshReq.Header.Set("Authorization", "Bearer "+tokenString)
	recorded = test.RunRequest(t, refreshAPI.MakeHandler(), validRefreshReq)
	recorded.CodeIs(200)
	recorded.ContentTypeIsJson()

	rToken := DecoderToken{}
	test.DecodeJsonPayload(recorded.Recorder, &rToken)
	refreshToken, err := jwt.Parse(rToken.Token, func(token *jwt.Token) (interface{}, error) {
		return key, nil
	})

	if err != nil {
		t.Errorf("Received refreshed token with wrong signature %s", err)
	}

	if refreshToken.Claims["testkey"].(string) != "testval" {
		t.Errorf("Received new token without payload")
	}

	// payload is accessible in request
	payloadAPI := rest.NewApi()
	payloadAPI.Use(authMiddleware)
	payloadAPI.SetApp(rest.AppSimple(func(w rest.ResponseWriter, r *rest.Request) {
		testval := r.Env["JWT_PAYLOAD"].(map[string]interface{})["testkey"].(string)
		w.WriteJson(map[string]string{"testkey": testval})
	}))

	payloadToken := jwt.New(jwt.GetSigningMethod("HS256"))
	payloadToken.Claims["id"] = "admin"
	payloadToken.Claims["exp"] = time.Now().Add(time.Hour).Unix()
	payloadToken.Claims["orig_iat"] = time.Now().Unix()
	payloadToken.Claims["testkey"] = "testval"
	payloadTokenString, _ := payloadToken.SignedString(key)

	payloadReq := test.MakeSimpleRequest("GET", "http://localhost/", nil)
	payloadReq.Header.Set("Authorization", "Bearer "+payloadTokenString)
	recorded = test.RunRequest(t, payloadAPI.MakeHandler(), payloadReq)
	recorded.CodeIs(200)
	recorded.ContentTypeIsJson()

	payload := map[string]string{}
	test.DecodeJsonPayload(recorded.Recorder, &payload)

//.........这里部分代码省略.........