From 0d482508fce00ab77cfb8f9f7584e13fc035aef4 Mon Sep 17 00:00:00 2001
From: Makis Maropoulos <makis@ideopod.com>
Date: Thu, 16 Jun 2016 20:08:59 +0300
Subject: [PATCH] Add ReadBufferSize and WriteBufferSize to the iris-ws,
 disable header copy also

---
 config/websocket.go | 18 ++++++++++++------
 http.go             | 15 ++++++---------
 iris.go             | 13 +++++++++++--
 websocket/server.go |  2 +-
 4 files changed, 30 insertions(+), 18 deletions(-)

diff --git a/config/websocket.go b/config/websocket.go
index cc606f76..68885c0f 100644
--- a/config/websocket.go
+++ b/config/websocket.go
@@ -40,17 +40,23 @@ type Websocket struct {
 	// Headers  the response headers before upgrader
 	// Default is empty
 	Headers map[string]string
+	// ReadBufferSize is the buffer size for the underline reader
+	ReadBufferSize int
+	// WriteBufferSize is the buffer size for the underline writer
+	WriteBufferSize int
 }
 
 // DefaultWebsocket returns the default config for iris-ws websocket package
 func DefaultWebsocket() *Websocket {
 	return &Websocket{
-		WriteTimeout:   DefaultWriteTimeout,
-		PongTimeout:    DefaultPongTimeout,
-		PingPeriod:     DefaultPingPeriod,
-		MaxMessageSize: DefaultMaxMessageSize,
-		Headers:        make(map[string]string, 0),
-		Endpoint:       "",
+		WriteTimeout:    DefaultWriteTimeout,
+		PongTimeout:     DefaultPongTimeout,
+		PingPeriod:      DefaultPingPeriod,
+		MaxMessageSize:  DefaultMaxMessageSize,
+		ReadBufferSize:  4096,
+		WriteBufferSize: 4096,
+		Headers:         make(map[string]string, 0),
+		Endpoint:        "",
 	}
 }
 
diff --git a/http.go b/http.go
index 04473499..15d5bd6b 100644
--- a/http.go
+++ b/http.go
@@ -311,7 +311,6 @@ func (s *Server) listen() error {
 	if s.started {
 		return errServerAlreadyStarted.Return()
 	}
-
 	listener, err := net.Listen("tcp4", s.Config.ListeningAddr)
 
 	if err != nil {
@@ -362,7 +361,7 @@ func (s *Server) serve(l net.Listener) error {
 
 // Open opens/starts/runs/listens (to) the server, listen tls if Cert && Key is registed, listenUNIX if Mode is registed, otherwise listen
 func (s *Server) Open() error {
-	if s.started {
+	if s.IsListening() {
 		return errServerAlreadyStarted.Return()
 	}
 
@@ -389,16 +388,14 @@ func (s *Server) Open() error {
 }
 
 // close closes the server
-func (s *Server) close() error {
-	if !s.started {
+func (s *Server) close() (err error) {
+	if !s.started || s.listener == nil {
 		return errServerIsClosed.Return()
 	}
+	s.started = false
+	err = s.listener.Close()
 
-	if s.listener != nil {
-		s.started = false
-		return s.listener.Close()
-	}
-	return nil
+	return
 }
 
 // errHandler returns na error with message: 'Passed argument is not func(*Context) neither an object which implements the iris.Handler with Serve(ctx *Context)
diff --git a/iris.go b/iris.go
index f15929c1..c48236c5 100644
--- a/iris.go
+++ b/iris.go
@@ -299,14 +299,24 @@ func (s *Framework) NoListen() *Server {
 	return s.HTTPServer
 }
 
+// CloseWithErr terminates the server and returns an error if any
+func CloseWithErr() error {
+	return Default.CloseWithErr()
+}
+
 //Close terminates the server and panic if error occurs
 func Close() {
 	Default.Close()
 }
 
+// CloseWithErr terminates the server and returns an error if any
+func (s *Framework) CloseWithErr() error {
+	return s.closeServer()
+}
+
 //Close terminates the server and panic if error occurs
 func (s *Framework) Close() {
-	s.Must(s.closeServer())
+	s.Must(s.CloseWithErr())
 }
 
 // OnError registers a custom http error handler
@@ -1143,7 +1153,6 @@ func StaticContent(reqPath string, contentType string, content []byte) RouteName
 func (api *muxAPI) StaticContent(reqPath string, cType string, content []byte) func(string) { // func(string) because we use that on websockets
 	modtime := time.Now()
 	modtimeStr := modtime.UTC().Format(config.TimeFormat)
-
 	h := func(ctx *Context) {
 		if t, err := time.Parse(config.TimeFormat, ctx.RequestHeader(ifModifiedSince)); err == nil && modtime.Before(t.Add(config.StaticCacheDuration)) {
 			ctx.Response.Header.Del(contentType)
diff --git a/websocket/server.go b/websocket/server.go
index 274312d0..49cc936d 100644
--- a/websocket/server.go
+++ b/websocket/server.go
@@ -73,7 +73,7 @@ func newServer(c *config.Websocket) *server {
 		onConnectionListeners: make([]ConnectionFunc, 0),
 	}
 
-	s.upgrader = websocket.New(s.handleConnection)
+	s.upgrader = websocket.Custom(s.handleConnection, s.config.ReadBufferSize, s.config.WriteBufferSize, false)
 	go s.serve() // start the server automatically
 	return s
 }