Golang context.Set函数代码示例

// 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 UserTFAPage(w http.ResponseWriter, req *http.Request) {
	args := handlers.GetArgs(req)

	u := quimby.NewUser(args.Vars["username"], quimby.UserDB(args.DB), quimby.UserTFA(handlers.TFA))
	if err := u.Fetch(); err != nil {
		context.Set(req, "error", err)
		return
	}

	qrData, err := u.UpdateTFA()
	if err != nil {
		context.Set(req, "error", err)
		return
	}

	if _, err := u.Save(); err != nil {
		context.Set(req, "error", err)
		return
	}

	qr := qrPage{
		userPage: userPage{
			User:  args.User.Username,
			Admin: handlers.Admin(args),
			Links: []link{
				{"quimby", "/"},
				{"admin", "/admin.html"},
			},
		},
		QR: template.HTMLAttr(base64.StdEncoding.EncodeToString(qrData)),
	}
	templates["qr-code.html"].template.ExecuteTemplate(w, "base", qr)
}

// SetUser is a HTTP middleware for setting the user and build information.
func SetUser(h http.Handler) http.Handler {
	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
		context.Set(req, csrf.TemplateTag, csrf.TemplateField(req)) // Set CSRF field everywhere
		context.Set(req, "Version", Version)
		context.Set(req, "Branch", Branch)
		context.Set(req, "user", &Account{})
		ses, err := store.Get(req, "evemu")
		if err != nil {
			log.Printf("Failed to get session cookie: %s", err.Error())
			h.ServeHTTP(rw, req)
			return
		}

		if v, ok := ses.Values["user"]; ok {
			if v.(int) == 0 {
				h.ServeHTTP(rw, req)
				return
			}
			//a, err := DBAccountFromID(v.(int))
			a, err := Layer.GetAccountFromID(v.(int)) // TODO: Unsafe, may panic
			if err != nil {
				log.Printf("Failed to get account from id: %s", err.Error())
				h.ServeHTTP(rw, req)
				return
			}
			context.Set(req, "user", a)
		}
		h.ServeHTTP(rw, req)
	})
}

func TestBuildURL(t *testing.T) {
	assert := assert.New(t)

	api := NewAPI(NewConfiguration())
	api.RegisterResourceHandler(TestResourceHandler{})
	api.RegisterResourceHandler(ComplexTestResourceHandler{})

	req, _ := http.NewRequest("GET", "https://example.com/api/v1/widgets", nil)
	gContext.Set(req, "version", "1")

	ctx := NewContextWithRouter(nil, req, api.(*muxAPI).router)

	url, _ := ctx.BuildURL("widgets", HandleCreate, nil)
	assert.Equal(url, "http://example.com/api/v1/widgets")

	url, _ = ctx.BuildURL("widgets", HandleRead, RouteVars{"resource_id": "111"})
	assert.Equal(url, "http://example.com/api/v1/widgets/111")

	url, _ = ctx.BuildPath("widgets", HandleRead, RouteVars{"resource_id": "111"})
	assert.Equal(url, "/api/v1/widgets/111")

	// Secure request should produce https URL
	req.TLS = &tls.ConnectionState{}
	url, _ = ctx.BuildURL("widgets", HandleRead, RouteVars{"resource_id": "222"})
	assert.Equal(url, "https://example.com/api/v1/widgets/222")

	// Make sure this works with another version number
	gContext.Set(req, "version", "2")
	url, _ = ctx.BuildURL("resources", HandleCreate, RouteVars{
		"company":  "acme",
		"category": "anvils"})
	assert.Equal(url, "https://example.com/api/v2/acme/anvils/resources")
}

func UpdateGadget(w http.ResponseWriter, req *http.Request) {
	args := GetArgs(req)
	var g quimby.Gadget
	dec := json.NewDecoder(req.Body)
	err := dec.Decode(&g)
	if err != nil {
		context.Set(req, "error", err)
		return // err
	}

	g.DB = args.DB
	g.Id = args.Vars["id"]
	err = g.Save()
	if err != nil {
		context.Set(req, "error", err)
		return // err
	}

	u, err := url.Parse(fmt.Sprintf("/api/gadgets/%s", g.Id))
	if err != nil {
		context.Set(req, "error", err)
		return //err
	}
	w.Header().Set("Location", u.String())
}

func UserChangePasswordPage(w http.ResponseWriter, req *http.Request) {
	args := handlers.GetArgs(req)
	u := quimby.NewUser(args.Vars["username"], quimby.UserDB(args.DB))
	if err := u.Fetch(); err != nil {
		context.Set(req, "error", err)
		return
	}
	if err := req.ParseForm(); err != nil {
		context.Set(req, "error", err)
		return
	}

	u.Password = req.PostFormValue("password")
	pw := req.PostFormValue("password_confirm")
	if pw != u.Password {
		context.Set(req, "error", ErrPasswordsDoNotMatch)
		return
	}

	if _, err := u.Save(); err != nil {
		context.Set(req, "error", ErrPasswordsDoNotMatch)
		return
	}

	w.Header().Set("Location", "/admin.html")
	w.WriteHeader(http.StatusFound)
}

func HandlerRegister(rw http.ResponseWriter, req *http.Request) {
	switch req.Method {
	case "POST":
		username := req.FormValue("username")
		password := req.FormValue("password")
		password2 := req.FormValue("password-c")
		if username == "" {
			context.Set(req, "error", fmt.Errorf("Username can't be blank"))
			DoTemplate("error.tmpl", rw, req)
			return
		}
		if password == "" || password == "" {
			context.Set(req, "error", fmt.Errorf("The passwords can't be blank"))
			DoTemplate("error.tmpl", rw, req)
			return
		}
		if password != password2 {
			context.Set(req, "error", fmt.Errorf("The passwords didn't match"))
			DoTemplate("error.tmpl", rw, req)
			return
		}
		err := Layer.NewAccount(username, password)
		if err != nil {
			context.Set(req, "error", err)
			DoTemplate("error.tmpl", rw, req)
			return
		}
		context.Set(req, "registered", true)
	}
	DoTemplate("register.tmpl", rw, req)
}

func (rc *ResourceController) CreateHandler(rw http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
	decoder := json.NewDecoder(r.Body)
	resource := models.NewStructForResourceName(rc.Name)
	err := decoder.Decode(resource)
	if err != nil {
		http.Error(rw, err.Error(), http.StatusInternalServerError)
	}

	c := Database.C(models.PluralizeLowerResourceName(rc.Name))
	i := bson.NewObjectId()
	reflect.ValueOf(resource).Elem().FieldByName("Id").SetString(i.Hex())
	err = c.Insert(resource)
	if err != nil {
		http.Error(rw, err.Error(), http.StatusInternalServerError)
	}

	context.Set(r, rc.Name, resource)
	context.Set(r, "Resource", rc.Name)
	context.Set(r, "Action", "create")

	rw.Header().Add("Location", responseURL(r, rc.Name, i.Hex()).String())
	rw.Header().Set("Content-Type", "application/json; charset=utf-8")
	rw.Header().Set("Access-Control-Allow-Origin", "*")
	rw.WriteHeader(http.StatusCreated)
	json.NewEncoder(rw).Encode(resource)
}

// GetContext wraps each request in a function which fills in the context for a given request.
// This includes setting the User and Session keys and values as necessary for use in later functions.
func GetContext(handler http.Handler) http.HandlerFunc {
	// Set the context here
	return func(w http.ResponseWriter, r *http.Request) {
		// Parse the request form
		err := r.ParseForm()
		if err != nil {
			http.Error(w, "Error parsing request", http.StatusInternalServerError)
		}
		// Set the context appropriately here.
		// Set the session
		session, _ := auth.Store.Get(r, "gophish")
		// Put the session in the context so that
		ctx.Set(r, "session", session)
		if id, ok := session.Values["id"]; ok {
			u, err := models.GetUser(id.(int64))
			if err != nil {
				ctx.Set(r, "user", nil)
			} else {
				ctx.Set(r, "user", u)
			}
		} else {
			ctx.Set(r, "user", nil)
		}
		handler.ServeHTTP(w, r)
		// Remove context contents
		ctx.Clear(r)
	}
}

// TODO: Do something when the user doesn't POST here
func HandlerLogin(rw http.ResponseWriter, req *http.Request) {
	switch req.Method {
	case "POST":
		username := req.FormValue("username")
		password := req.FormValue("password")
		a, err := Layer.Login(username, password)
		if err == sql.ErrNoRows {
			context.Set(req, "error", fmt.Errorf("db: User doesn't exist"))
			DoTemplate("error.tmpl", rw, req)
			return
		} else if err != nil {
			context.Set(req, "error", err)
			context.Set(req, "stack", string(debug.Stack()))
			DoTemplate("500.tmpl", rw, req)
			return
		}

		context.Set(req, "user", *a)
		ses, err := store.Get(req, "evemu")
		if err != nil {
			context.Set(req, "error", err)
			DoTemplate("500.tmpl", rw, req)
			return
		}

		ses.Values["user"] = a.ID
		ses.Save(req, rw)

		http.Redirect(rw, req, "/", http.StatusFound)
	}
}

// authenticateLoader is used to authenticate requests that are made to the
// loader API endpoints. Rather than operate on GPG signatures, the
// authentication instead uses the submitted loader key
func authenticateLoader(pass handler) handler {
	return func(w http.ResponseWriter, r *http.Request) {
		var (
			loaderid float64
			err      error
		)
		opid := getOpID(r)
		context.Set(r, opID, opid)
		lkey := r.Header.Get("X-LOADERKEY")
		if lkey == "" {
			resource := cljs.New(fmt.Sprintf("%s%s", ctx.Server.Host, r.URL.String()))
			resource.SetError(cljs.Error{Code: fmt.Sprintf("%.0f", opid), Message: "X-LOADERKEY header not found"})
			respond(http.StatusUnauthorized, resource, w, r)
			return
		}
		// Do a sanity check here on the submitted loader string before
		// we attempt the authentication
		err = mig.ValidateLoaderKey(lkey)
		if err == nil {
			loaderid, err = ctx.DB.GetLoaderEntryID(lkey)
		}
		if err != nil {
			resource := cljs.New(fmt.Sprintf("%s%s", ctx.Server.Host, r.URL.String()))
			resource.SetError(cljs.Error{Code: fmt.Sprintf("%.0f", opid), Message: fmt.Sprintf("Loader authorization failed")})
			respond(http.StatusUnauthorized, resource, w, r)
			return
		}
		context.Set(r, loaderID, loaderid)
		// accept request
		pass(w, r)
	}
}

// NewContext returns a RequestContext populated with parameters from the request path and
// query string.
func NewContext(parent context.Context, req *http.Request) RequestContext {
	if parent == nil {
		parent = context.Background()
	}

	for key, value := range req.URL.Query() {
		var val interface{}
		val = value

		// Query string values are slices (e.g. ?foo=bar,baz,qux yields
		// [bar, baz, qux] for foo), but we unbox single values (e.g. ?foo=bar
		// yields bar for foo).
		if len(value) == 1 {
			val = value[0]
		}

		gcontext.Set(req, key, val)
	}

	for key, value := range mux.Vars(req) {
		gcontext.Set(req, key, value)
	}

	// TODO: Keys can potentially be overwritten if the request path has
	// parameters with the same name as query string values. Figure out a
	// better way to handle this.

	return &gorillaRequestContext{parent, req, []string{}}
}

func (sm *ServiceManager) authenticateRequest(r *http.Request) (username string, err error) {
	isAdmin := false

	token, err := auth.GetTokenFromRequest(r, []byte(sm.cfg.Auth.Token.SigningKey))
	if err == nil {
		sm.log.Tracef("Token validated: %#v\n", *token)

		username = token.Claims["sub"].(string)

		if _, ok := token.Claims["admin"]; ok {
			isAdmin, _ = token.Claims["admin"].(bool)
		}
	} else if strings.Contains(err.Error(), "expired") {
		return
	} else {
		sm.log.Debugf("Skipping token auth: %s\n", err)

		var cacheHit bool
		if username, cacheHit, err = sm.authenticator.AuthenticateRequest(r); err != nil {
			return
		}
		sm.log.Tracef("Cache hit (%s): %v\n", username, cacheHit)

		isAdmin = sm.localAuthGroups.UserHasGroupMembership(username, "admin")
	}

	sm.log.Tracef("Setting request context: IsAdmin=%v\n", isAdmin)
	context.Set(r, handlers.IsAdmin, isAdmin)

	sm.log.Tracef("Setting request context: Username=%s\n", username)
	context.Set(r, handlers.Username, username)
	return
}

// ServeHTTP as per the negroni.Handler interface
func (m *MixedAuthMiddleware) ServeHTTP(w http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
	// HTTPS redirection
	err := util.NewSecure(m.service.GetConfig().IsDevelopment).Process(w, r)
	if err != nil {
		return
	}

	account, user, err := getMixedCredentialsFromRequest(r, m.service)

	if err != nil {
		// For security reasons, return a generic error message
		response.UnauthorizedError(w, ErrAccountOrUserAuthenticationRequired.Error())
		return
	}

	if account != nil {
		context.Set(r, AuthenticatedAccountKey, account)
	}

	if user != nil {
		context.Set(r, AuthenticatedUserKey, user)
	}

	next(w, r)
}

func (k *AuthKey) ProcessRequest(w http.ResponseWriter, r *http.Request, configuration interface{}) (error, int) {
	var thisConfig AuthKeyConfiguration
	thisConfig = configuration.(AuthKeyConfiguration)

	authHeaderValue := r.Header.Get(thisConfig.Auth.AuthHeaderName)
	if authHeaderValue == "" {
		// No header value, fail
		log.WithFields(logrus.Fields{
			"path":   r.URL.Path,
			"origin": r.RemoteAddr,
		}).Info("Attempted access with malformed header, no auth header found.")

		return errors.New("Authorization field missing"), 400
	}

	// Check if API key valid
	thisSessionState, keyExists := k.TykMiddleware.CheckSessionAndIdentityForValidKey(authHeaderValue)
	if !keyExists {
		log.WithFields(logrus.Fields{
			"path":   r.URL.Path,
			"origin": r.RemoteAddr,
			"key":    authHeaderValue,
		}).Info("Attempted access with non-existent key.")

		return errors.New("Key not authorised"), 403
	}

	// Set session state on context, we will need it later
	context.Set(r, SessionData, thisSessionState)
	context.Set(r, AuthHeaderValue, authHeaderValue)

	return nil, 200
}