From 1533bb555e1423d81d7accea9a545f0f23880935 Mon Sep 17 00:00:00 2001
From: Daniel Cormier <dcormier@users.noreply.github.com>
Date: Tue, 29 Oct 2019 10:10:10 -0400
Subject: [PATCH] Set the default domain for cookies properly (#187)

Fixes #186.
---
 samlsp/middleware_test.go | 28 ++++++++++++++++++++++++++++
 samlsp/samlsp.go          |  7 ++++++-
 2 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/samlsp/middleware_test.go b/samlsp/middleware_test.go
index 7ce2a6df..6d7dd47d 100644
--- a/samlsp/middleware_test.go
+++ b/samlsp/middleware_test.go
@@ -8,6 +8,7 @@ import (
 	"encoding/base64"
 	"encoding/xml"
 	"io/ioutil"
+	"net"
 	"net/http"
 	"net/http/httptest"
 	"net/url"
@@ -405,6 +406,33 @@ func (test *MiddlewareTest) TestCanParseResponse(c *C) {
 	})
 }
 
+func (test *MiddlewareTest) TestDefaultCookieDomainIPv4(c *C) {
+	ipv4Loopback := net.IP{127, 0, 0, 1}
+	mw, err := New(Options{
+		URL:         mustParseURL("https://" + net.JoinHostPort(ipv4Loopback.String(), "54321")),
+		Key:         test.Key,
+		Certificate: test.Certificate,
+		IDPMetadata: &saml.EntityDescriptor{},
+	})
+	c.Assert(err, IsNil)
+
+	cookieStore := mw.ClientToken.(*ClientCookies)
+	c.Assert(cookieStore.Domain, Equals, ipv4Loopback.String(), Commentf("Cookie domain must not contain a port or the cookie cannot be set properly"))
+}
+
+func (test *MiddlewareTest) TestDefaultCookieDomainIPv6(c *C) {
+	mw, err := New(Options{
+		URL:         mustParseURL("https://" + net.JoinHostPort(net.IPv6loopback.String(), "54321")),
+		Key:         test.Key,
+		Certificate: test.Certificate,
+		IDPMetadata: &saml.EntityDescriptor{},
+	})
+	c.Assert(err, IsNil)
+
+	cookieStore := mw.ClientToken.(*ClientCookies)
+	c.Assert(cookieStore.Domain, Equals, net.IPv6loopback.String(), Commentf("Cookie domain must not contain a port or the cookie cannot be set properly"))
+}
+
 func (test *MiddlewareTest) TestRejectsInvalidRelayState(c *C) {
 	v := &url.Values{}
 	v.Set("SAMLResponse", base64.StdEncoding.EncodeToString([]byte(test.SamlResponse)))
diff --git a/samlsp/samlsp.go b/samlsp/samlsp.go
index 52e82f9a..4f7fac38 100644
--- a/samlsp/samlsp.go
+++ b/samlsp/samlsp.go
@@ -8,6 +8,7 @@ import (
 	"encoding/xml"
 	"fmt"
 	"io/ioutil"
+	"net"
 	"net/http"
 	"net/url"
 	"time"
@@ -77,7 +78,11 @@ func New(opts Options) (*Middleware, error) {
 			if opts.CookieDomain != "" {
 				return opts.CookieDomain
 			}
-			return opts.URL.Host
+			host, _, err := net.SplitHostPort(opts.URL.Host)
+			if err != nil {
+				return opts.URL.Host
+			}
+			return host
 		}(),
 		Secure: opts.CookieSecure,
 	}