From a7afffe70082cb28bc73525a8b2b7c50942fde61 Mon Sep 17 00:00:00 2001
From: Lucas Clemente <lucas@clemente.io>
Date: Tue, 31 May 2016 15:38:21 +0200
Subject: [PATCH] add a SetQuicHeaders function to h2quic.Server

fixes #156
---
 h2quic/server.go         | 28 ++++++++++++++++++++++++++
 h2quic/server_test.go    | 43 ++++++++++++++++++++++++++++++++++++++++
 protocol/version.go      | 11 ++++++++++
 protocol/version_test.go |  4 ++++
 4 files changed, 86 insertions(+)

diff --git a/h2quic/server.go b/h2quic/server.go
index 72f746c38..eeb998745 100644
--- a/h2quic/server.go
+++ b/h2quic/server.go
@@ -3,6 +3,7 @@ package h2quic
 import (
 	"crypto/tls"
 	"errors"
+	"fmt"
 	"io/ioutil"
 	"net"
 	"net/http"
@@ -26,6 +27,8 @@ type Server struct {
 
 	// Private flag for demo, do not use
 	CloseAfterFirstRequest bool
+
+	port int
 }
 
 // ListenAndServe listens on the UDP address s.Addr and calls s.Handler to handle HTTP/2 requests on incoming connections.
@@ -142,11 +145,36 @@ func (s *Server) handleRequest(session streamCreator, headerStream utils.Stream,
 	return nil
 }
 
+// Close the server
 func (s *Server) Close() error {
 	// TODO: implement
 	return nil
 }
 
+// SetQuicHeaders can be used to set the proper headers that announce that this server supports QUIC.
+// The values that are set depend on the port information from s.Server.Addr, and currently look like this (if Addr has port 443):
+//  Alternate-Protocol: 443:quic
+//  Alt-Svc: quic=":443"; ma=2592000; v="33,32,31,30"
+func (s *Server) SetQuicHeaders(hdr http.Header) error {
+	if s.port == 0 {
+		// Extract port from s.Server.Addr
+		_, portStr, err := net.SplitHostPort(s.Server.Addr)
+		if err != nil {
+			return err
+		}
+		port, err := net.LookupPort("tcp", portStr)
+		if err != nil {
+			return err
+		}
+		s.port = port
+	}
+
+	hdr.Add("Alternate-Protocol", fmt.Sprintf("%d:quic", s.port))
+	hdr.Add("Alt-Svc", fmt.Sprintf(`quic=":%d"; ma=2592000; v="%s"`, s.port, protocol.SupportedVersionsAsString))
+
+	return nil
+}
+
 // ListenAndServeQUIC listens on the UDP network address addr and calls the
 // handler for HTTP/2 requests on incoming conections. http.DefaultServeMux is
 // used when handler is nil.
diff --git a/h2quic/server_test.go b/h2quic/server_test.go
index cc1bc7641..9fcd13bc8 100644
--- a/h2quic/server_test.go
+++ b/h2quic/server_test.go
@@ -157,4 +157,47 @@ var _ = Describe("H2 server", func() {
 			Server{}.Serve(nil)
 		}).To(Panic())
 	})
+
+	Context("setting http headers", func() {
+		expected := http.Header{
+			"Alt-Svc":            {`quic=":443"; ma=2592000; v="33,32,31,30"`},
+			"Alternate-Protocol": {`443:quic`},
+		}
+
+		It("sets proper headers with numeric port", func() {
+			s.Server.Addr = ":443"
+			hdr := http.Header{}
+			err := s.SetQuicHeaders(hdr)
+			Expect(err).NotTo(HaveOccurred())
+			Expect(hdr).To(Equal(expected))
+		})
+
+		It("sets proper headers with full addr", func() {
+			s.Server.Addr = "127.0.0.1:443"
+			hdr := http.Header{}
+			err := s.SetQuicHeaders(hdr)
+			Expect(err).NotTo(HaveOccurred())
+			Expect(hdr).To(Equal(expected))
+		})
+
+		It("sets proper headers with string port", func() {
+			s.Server.Addr = ":https"
+			hdr := http.Header{}
+			err := s.SetQuicHeaders(hdr)
+			Expect(err).NotTo(HaveOccurred())
+			Expect(hdr).To(Equal(expected))
+		})
+
+		It("works multiple times", func() {
+			s.Server.Addr = ":https"
+			hdr := http.Header{}
+			err := s.SetQuicHeaders(hdr)
+			Expect(err).NotTo(HaveOccurred())
+			Expect(hdr).To(Equal(expected))
+			hdr = http.Header{}
+			err = s.SetQuicHeaders(hdr)
+			Expect(err).NotTo(HaveOccurred())
+			Expect(hdr).To(Equal(expected))
+		})
+	})
 })
diff --git a/protocol/version.go b/protocol/version.go
index 21ee6795d..12e2842e7 100644
--- a/protocol/version.go
+++ b/protocol/version.go
@@ -3,6 +3,7 @@ package protocol
 import (
 	"bytes"
 	"encoding/binary"
+	"strconv"
 )
 
 // VersionNumber is a version number as int
@@ -16,6 +17,9 @@ var SupportedVersions = []VersionNumber{
 // SupportedVersionsAsTags is needed for the SHLO crypto message
 var SupportedVersionsAsTags []byte
 
+// SupportedVersionsAsString is needed for the Alt-Scv HTTP header
+var SupportedVersionsAsString string
+
 // VersionNumberToTag maps version numbers ('32') to tags ('Q032')
 func VersionNumberToTag(vn VersionNumber) uint32 {
 	v := uint32(vn)
@@ -45,4 +49,11 @@ func init() {
 		b.Write(s)
 	}
 	SupportedVersionsAsTags = b.Bytes()
+
+	for i := len(SupportedVersions) - 1; i >= 0; i-- {
+		SupportedVersionsAsString += strconv.Itoa(int(SupportedVersions[i]))
+		if i != 0 {
+			SupportedVersionsAsString += ","
+		}
+	}
 }
diff --git a/protocol/version_test.go b/protocol/version_test.go
index 7812917b9..897a4221a 100644
--- a/protocol/version_test.go
+++ b/protocol/version_test.go
@@ -22,6 +22,10 @@ var _ = Describe("Version", func() {
 		Expect(protocol.SupportedVersionsAsTags).To(Equal([]byte("Q030Q031Q032Q033")))
 	})
 
+	It("has proper version list", func() {
+		Expect(protocol.SupportedVersionsAsString).To(Equal("33,32,31,30"))
+	})
+
 	It("recognizes supported versions", func() {
 		Expect(protocol.IsSupportedVersion(0)).To(BeFalse())
 		Expect(protocol.IsSupportedVersion(protocol.SupportedVersions[0])).To(BeTrue())