Golang context.GetOk函数代码示例

func TestContextBindAndValidate(t *testing.T) {
	spec, api := petstore.NewAPI(t)
	ctx := NewContext(spec, api, nil)
	ctx.router = DefaultRouter(spec, ctx.api)

	request, _ := http.NewRequest("POST", "/pets", nil)
	request.Header.Add("Accept", "*/*")
	request.Header.Add("content-type", "text/html")

	v, ok := context.GetOk(request, ctxBoundParams)
	assert.False(t, ok)
	assert.Nil(t, v)

	ri, _ := ctx.RouteInfo(request)
	data, result := ctx.BindAndValidate(request, ri) // this requires a much more thorough test
	assert.NotNil(t, data)
	assert.NotNil(t, result)

	v, ok = context.GetOk(request, ctxBoundParams)
	assert.True(t, ok)
	assert.NotNil(t, v)

	dd, rr := ctx.BindAndValidate(request, ri)
	assert.Equal(t, data, dd)
	assert.Equal(t, result, rr)
}

func TestContextInvalidResponseFormat(t *testing.T) {
	ct := "application/x-yaml"
	other := "application/sgml"
	spec, api := petstore.NewAPI(t)
	ctx := NewContext(spec, api, nil)
	ctx.router = DefaultRouter(spec, ctx.api)

	request, _ := http.NewRequest("GET", "http://localhost:8080", nil)
	request.Header.Set(httpkit.HeaderAccept, ct)

	// check there's nothing there
	cached, ok := context.GetOk(request, ctxResponseFormat)
	assert.False(t, ok)
	assert.Empty(t, cached)

	// trigger the parse
	mt := ctx.ResponseFormat(request, []string{other})
	assert.Empty(t, mt)

	// check it was cached
	cached, ok = context.GetOk(request, ctxResponseFormat)
	assert.False(t, ok)
	assert.Empty(t, cached)

	// check if the cast works and fetch from cache too
	mt = ctx.ResponseFormat(request, []string{other})
	assert.Empty(t, mt)
}

func TestContextValidContentType(t *testing.T) {
	ct := "application/json"
	ctx := NewContext(nil, nil, nil)

	request, _ := http.NewRequest("GET", "http://localhost:8080", nil)
	request.Header.Set(httpkit.HeaderContentType, ct)

	// check there's nothing there
	_, ok := context.GetOk(request, ctxContentType)
	assert.False(t, ok)

	// trigger the parse
	mt, _, err := ctx.ContentType(request)
	assert.NoError(t, err)
	assert.Equal(t, ct, mt)

	// check it was cached
	_, ok = context.GetOk(request, ctxContentType)
	assert.True(t, ok)

	// check if the cast works and fetch from cache too
	mt, _, err = ctx.ContentType(request)
	assert.NoError(t, err)
	assert.Equal(t, ct, mt)
}

// Tests that the context is cleared or not cleared properly depending on
// the configuration of the router
func TestKeepContext(t *testing.T) {
	func1 := func(w http.ResponseWriter, r *http.Request) {}

	r := NewRouter()
	r.HandleFunc("/", func1).Name("func1")

	req, _ := http.NewRequest("GET", "http://localhost/", nil)
	context.Set(req, "t", 1)

	res := new(http.ResponseWriter)
	r.ServeHTTP(*res, req)

	if _, ok := context.GetOk(req, "t"); ok {
		t.Error("Context should have been cleared at end of request")
	}

	r.KeepContext = true

	req, _ = http.NewRequest("GET", "http://localhost/", nil)
	context.Set(req, "t", 1)

	r.ServeHTTP(*res, req)
	if _, ok := context.GetOk(req, "t"); !ok {
		t.Error("Context should NOT have been cleared at end of request")
	}

}

func currentServicePipeID(r *http.Request) (string, string) {
	var serviceID, pipeID string
	if v, ok := context.GetOk(r, serviceIDKey); ok {
		serviceID = v.(string)
	}
	if v, ok := context.GetOk(r, pipeIDKey); ok {
		pipeID = v.(string)
	}
	return serviceID, pipeID
}

func GetFeatures(r *http.Request) map[string]interface{} {
	result := map[string]interface{}{}

	if features, ok := context.GetOk(r, "features"); ok {
		result["features"] = features
	}

	if identifier, ok := context.GetOk(r, "api_key"); ok {
		result["identifier"] = identifier
	}

	result["timestamp"] = time.Now().UTC().Format(time.RFC3339Nano)

	return result
}

// GetCallID gets the current call ID (if any) from the request context.
func GetCallID(r *http.Request) (int64, bool) {
	if v, present := context.GetOk(r, callID); present {
		id, ok := v.(int64)
		return id, ok
	}
	return 0, false
}

func contextGet(r *http.Request, key string) (interface{}, error) {
	if val, ok := context.GetOk(r, key); ok {
		return val, nil
	}

	return nil, errors.Errorf("no value exists in the context for key %q", key)
}

func userContext(r *http.Request) (*account.Account, bool) {
	if user, ok := context.GetOk(r, userAccountKey); ok {
		u, ok := user.(*account.Account)
		return u, ok
	}
	return nil, false
}

func HandleSnapshots(w http.ResponseWriter, r *http.Request) {
	database, ok := context.GetOk(r, "DB")
	if !ok {
		RespondWithError(w, r, errors.New("Could'nt obtain database"), http.StatusInternalServerError)
		return
	}
	db := database.(*mgo.Database)
	params := r.URL.Query()

	d := time.Now()
	if date, ok := params["date"]; ok {
		if t, err := time.Parse(timeLayout, date[0]); err == nil {
			d = t
		}
	}

	snapshot, err := hubspider.FindSnapshotByTime(db, d)
	if err != nil {
		if err == mgo.ErrNotFound {
			RespondWithError(w, r, err, http.StatusNotFound)
			return
		}
		RespondWithError(w, r, err, http.StatusInternalServerError)
		return
	}
	snapshot.ServeJSON(w, r)
}

func (h *DBConn) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	logger.Info.Println("Setting up the database connection")
	_, ok := context.GetOk(r, consts.DB_KEY)

	if ok {
		errJson := `{"Error":"Internal Server Error"}`
		logger.Error.Println("DBConnector middleware error. DB_KEY already set.")
		http.Error(w, errJson, http.StatusInternalServerError)
		return
	}

	//https://github.com/mattn/go-sqlite3/blob/master/_example/simple/simple.go
	db, err := sql.Open("sqlite3", "./middleware-test.db")

	if err != nil {
		errJson := `{"Error":"Internal Server Error"}`
		logger.Error.Println("DBConnector database connection problem. Check Logs.")
		http.Error(w, errJson, http.StatusInternalServerError)
		return
	}

	context.Set(r, consts.DB_KEY, db)

	//Close database connection once middleware chain is complete
	defer db.Close()

	if h.next != nil {
		h.next.ServeHTTP(w, r)
	}
}

func main() {
	middle := interpose.New()

	// Create a middleware that yields a global counter that increments until
	// the server is shut down. Note that this would actually require a mutex
	// or a channel to be safe for concurrent use. Therefore, this example is
	// unsafe.
	middle.Use(context.ClearHandler)
	middle.Use(func() func(http.Handler) http.Handler {
		c := 0

		return func(next http.Handler) http.Handler {
			return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
				c++
				context.Set(req, CountKey, c)
				next.ServeHTTP(w, req)
			})
		}
	}())

	// Apply the router.
	router := mux.NewRouter()
	router.HandleFunc("/test/{user}", func(w http.ResponseWriter, req *http.Request) {
		c, ok := context.GetOk(req, CountKey)
		if !ok {
			fmt.Println("Context not ok")
		}

		fmt.Fprintf(w, "Hi %s, this is visit #%d to the site since the server was last rebooted.", mux.Vars(req)["user"], c)

	})
	middle.UseHandler(router)

	http.ListenAndServe(":3001", middle)
}

// GetPost() returns binded Post from POST data
func GetPost(r *http.Request) (Post, error) {
	rv, ok := context.GetOk(r, "post")
	if !ok {
		return Post{}, errors.New("context not set")
	}
	return rv.(Post), nil
}

//
//	HandleWith() constructs a handler and inserts it into
//	the mux router.
//
//	It does the following to each request:
//
//	Request Body JSON:
//  - If blank, fills body with {}
//  - Validates based on the schema provided
//		(default: dozy.PresetJsonSchemaBlank)
//
//	Query Strings:
//  - Converts all query strings, except authtoken, to a
//		map[string][]string
//
//	Rate Limiting:
//  - Handles rate limitng on requests with the authtoken in a
//	  URL query string called "authtoken"
//		- Because it only rate handles requests with "authtoken,"
//			it's especially important to use the .Query() function
//			provided by gorilla/mux as much as possible
//
//	Other:
//  - Writes the data specified by the user's handler function,
//		ensuring that only allowed HTTP response codes are used
//		and that no response is sent on error
//  - Sets the "Content-Type" header to "application/json"
//	- Manages the LOCATION header on a dozy.StatusPostOk return
//    value
//	- Ensures that Content-Type from client is correct
//
//	This method **must** be used with a dozy.ServeMux()-handled
//  mux router. dozy.ServeMux() attaches some neccesary data
//  values to the request, which are needed for the proper method
//  of the HandleWith() and the HandleSpecialCost() functions. However,
//  you can use a non dozy.HandleWith() method with the ServeMux()
//  method, only loosing the functionality added with the HandleWith()
//  method itself.
//
func (h handlerBuilder) HandleWith(method, path string, userHandler UserHandler) {
	method = strings.ToUpper(method)
	if (method == "GET" || method == "DELETE") && h.bodySchema != PresetJsonSchemaBlank {
		panic("Method type " + method + " should not have a body set!")
	}

	builtHandler := func(rw http.ResponseWriter, req *http.Request) {
		var cost uint

		if h.customCost {
			rawSettings, rawSettingsOk := context.GetOk(req, dozySettingsKey("settings"))
			settings, settingsOk := rawSettings.(*Settings)
			if !rawSettingsOk {
				panic("rawSettings is *not* okay!")
			}
			if !settingsOk {
				panic("settings is *not* okay!")
			}

			cost = settings.DefaultCost
		} else {
			cost = h.cost
		}

		handle(h.bodySchema, userHandler, cost, rw, req)
	}

	route := h.muxRouter.HandleFunc(path, builtHandler).Methods(method)
	if len(h.queries) != 0 {
		route.Queries(h.queries...)
	}
}

// Value returns Gorilla's context package's value for this Context's request
// and key. It delegates to the parent Context if there is no such value.
func (ctx *ctxWrapper) Value(key interface{}) interface{} {
	if val, ok := gorillactx.GetOk(ctx.req, key); ok {
		return val
	}

	return ctx.Context.Value(key)
}