From f5b1fb44bb819bd8059a84998c3a902bd01d92b3 Mon Sep 17 00:00:00 2001
From: Damon Zhao <yijun.zhao@outlook.com>
Date: Thu, 27 Aug 2015 16:04:50 +0800
Subject: [PATCH] Add SetCookie and GetCookie method for Context

---
 context.go      | 59 +++++++++++++++++++++++++++++++++++++++++++++++++
 context_test.go | 17 ++++++++++++++
 2 files changed, 76 insertions(+)

diff --git a/context.go b/context.go
index 20a20e33..c4fb96f8 100644
--- a/context.go
+++ b/context.go
@@ -9,6 +9,7 @@ import (
 	"io"
 	"math"
 	"net/http"
+	"net/url"
 	"strings"
 	"time"
 
@@ -321,6 +322,64 @@ func (c *Context) Header(key, value string) {
 	}
 }
 
+func (c *Context) SetCookie(name string, value string, options ...interface{}) {
+	cookie := http.Cookie{}
+	cookie.Name = name
+	cookie.Value = url.QueryEscape(value)
+
+	if len(options) > 0 {
+		switch v := options[0].(type) {
+		case int:
+			cookie.MaxAge = v
+		case int64:
+			cookie.MaxAge = int(v)
+		case int32:
+			cookie.MaxAge = int(v)
+		}
+	}
+
+	cookie.Path = "/"
+	if len(options) > 1 {
+		if v, ok := options[1].(string); ok && len(v) > 0 {
+			cookie.Path = v
+		}
+	}
+
+	if len(options) > 2 {
+		if v, ok := options[2].(string); ok && len(v) > 0 {
+			cookie.Domain = v
+		}
+	}
+
+	if len(options) > 3 {
+		switch v := options[3].(type) {
+		case bool:
+			cookie.Secure = v
+		default:
+			if options[3] != nil {
+				cookie.Secure = true
+			}
+		}
+	}
+
+	if len(options) > 4 {
+		if v, ok := options[4].(bool); ok && v {
+			cookie.HttpOnly = true
+		}
+	}
+
+	c.Writer.Header().Add("Set-Cookie", cookie.String())
+}
+
+func (c *Context) GetCookie(name string) string {
+	cookie, err := c.Request.Cookie(name)
+	if err != nil {
+		return ""
+	}
+	val, _ := url.QueryUnescape(cookie.Value)
+	return val
+}
+
 func (c *Context) Render(code int, r render.Render) {
 	c.writermem.WriteHeader(code)
 	if err := r.Render(c.Writer); err != nil {
diff --git a/context_test.go b/context_test.go
index d95125ae..5697a04d 100644
--- a/context_test.go
+++ b/context_test.go
@@ -246,6 +246,23 @@ func TestContextPostFormMultipart(t *testing.T) {
 	assert.Equal(t, c.PostForm("bar"), "foo")
 }
 
+func TestContextSetCookie(t *testing.T) {
+	c, w, _ := createTestContext()
+	c.SetCookie("user", "gin", 1, "/", "localhost", true, true)
+	c.SetCookie("user", "gin", int32(1), "/", "localhost", 1)
+	c.SetCookie("user", "gin", int64(1))
+
+	c.Request, _ = http.NewRequest("GET", "/set", nil)
+	assert.Equal(t, c.GetCookie("Set-Cookie"), "user=gin; Path=/; Domain=localhost; Max-Age=1; HttpOnly; Secure")
+}
+
+func TestContextGetCookie(t *testing.T) {
+	c, w, _ := createTestContext()
+	c.Request, _ = http.NewRequest("GET", "/get", nil)
+	c.Request.Header.Set("Cookie", "user=gin")
+	assert.Equal(t, w.Body.String(), "gin")
+}
+
 // Tests that the response is serialized as JSON
 // and Content-Type is set to application/json
 func TestContextRenderJSON(t *testing.T) {