From 12a9deb1a64108058d8ccd75705f95471f458cd0 Mon Sep 17 00:00:00 2001
From: Adnan Maolood <me@adnano.co>
Date: Tue, 27 Oct 2020 19:16:55 -0400
Subject: [PATCH] Make (*Response).Body an io.ReadCloser

---
 client.go          |  4 +---
 examples/client.go |  8 ++++++-
 response.go        | 53 +++++++++++++++++++++++++++++++++-------------
 3 files changed, 46 insertions(+), 19 deletions(-)

diff --git a/client.go b/client.go
index dac0329..03242e5 100644
--- a/client.go
+++ b/client.go
@@ -66,7 +66,6 @@ func (c *Client) Send(req *Request) (*Response, error) {
 	if err != nil {
 		return nil, err
 	}
-	defer conn.Close()
 
 	// Write the request
 	w := bufio.NewWriter(conn)
@@ -77,8 +76,7 @@ func (c *Client) Send(req *Request) (*Response, error) {
 
 	// Read the response
 	resp := &Response{}
-	r := bufio.NewReader(conn)
-	if err := resp.read(r); err != nil {
+	if err := resp.read(conn); err != nil {
 		return nil, err
 	}
 	// Store connection information
diff --git a/examples/client.go b/examples/client.go
index 659e783..fa88dcf 100644
--- a/examples/client.go
+++ b/examples/client.go
@@ -7,6 +7,7 @@ import (
 	"crypto/tls"
 	"crypto/x509"
 	"fmt"
+	"io/ioutil"
 	"net/url"
 	"os"
 	"time"
@@ -80,7 +81,12 @@ func sendRequest(req *gmi.Request) error {
 		req.URL.RawQuery = url.QueryEscape(scanner.Text())
 		return sendRequest(req)
 	case gmi.StatusClassSuccess:
-		fmt.Print(string(resp.Body))
+		defer resp.Body.Close()
+		body, err := ioutil.ReadAll(resp.Body)
+		if err != nil {
+			return err
+		}
+		fmt.Print(string(body))
 		return nil
 	case gmi.StatusClassRedirect:
 		fmt.Println("Redirecting to", resp.Meta)
diff --git a/response.go b/response.go
index 40d9046..4db127c 100644
--- a/response.go
+++ b/response.go
@@ -3,7 +3,7 @@ package gemini
 import (
 	"bufio"
 	"crypto/tls"
-	"io/ioutil"
+	"io"
 	"strconv"
 )
 
@@ -18,19 +18,20 @@ type Response struct {
 	// Meta should not be longer than 1024 bytes.
 	Meta string
 
-	// Body contains the response body.
-	Body []byte
+	// Body contains the response body for successful responses.
+	Body io.ReadCloser
 
 	// TLS contains information about the TLS connection on which the response
 	// was received.
 	TLS tls.ConnectionState
 }
 
-// read reads a Gemini response from the provided buffered reader.
-func (resp *Response) read(r *bufio.Reader) error {
+// read reads a Gemini response from the provided io.ReadCloser.
+func (resp *Response) read(rc io.ReadCloser) error {
+	br := bufio.NewReader(rc)
 	// Read the status
 	statusB := make([]byte, 2)
-	if _, err := r.Read(statusB); err != nil {
+	if _, err := br.Read(statusB); err != nil {
 		return err
 	}
 	status, err := strconv.Atoi(string(statusB))
@@ -47,14 +48,14 @@ func (resp *Response) read(r *bufio.Reader) error {
 	}
 
 	// Read one space
-	if b, err := r.ReadByte(); err != nil {
+	if b, err := br.ReadByte(); err != nil {
 		return err
 	} else if b != ' ' {
 		return ErrInvalidResponse
 	}
 
 	// Read the meta
-	meta, err := r.ReadString('\r')
+	meta, err := br.ReadString('\r')
 	if err != nil {
 		return err
 	}
@@ -67,19 +68,41 @@ func (resp *Response) read(r *bufio.Reader) error {
 	resp.Meta = meta
 
 	// Read terminating newline
-	if b, err := r.ReadByte(); err != nil {
+	if b, err := br.ReadByte(); err != nil {
 		return err
 	} else if b != '\n' {
 		return ErrInvalidResponse
 	}
 
-	// Read response body
 	if resp.Status.Class() == StatusClassSuccess {
-		var err error
-		resp.Body, err = ioutil.ReadAll(r)
-		if err != nil {
-			return err
-		}
+		resp.Body = newReadCloserBody(br, rc)
 	}
 	return nil
 }
+
+type readCloserBody struct {
+	br *bufio.Reader // used until empty
+	io.ReadCloser
+}
+
+func newReadCloserBody(br *bufio.Reader, rc io.ReadCloser) io.ReadCloser {
+	body := &readCloserBody{ReadCloser: rc}
+	if br.Buffered() != 0 {
+		body.br = br
+	}
+	return body
+}
+
+func (b *readCloserBody) Read(p []byte) (n int, err error) {
+	if b.br != nil {
+		if n := b.br.Buffered(); len(p) > n {
+			p = p[:n]
+		}
+		n, err = b.br.Read(p)
+		if b.br.Buffered() == 0 {
+			b.br = nil
+		}
+		return n, err
+	}
+	return b.ReadCloser.Read(p)
+}