Skip to content

Instantly share code, notes, and snippets.

@tlopo
Last active November 25, 2019 10:10
Show Gist options
  • Save tlopo/2561cf857237188aa563ebd4d9c6e5a4 to your computer and use it in GitHub Desktop.
Save tlopo/2561cf857237188aa563ebd4d9c6e5a4 to your computer and use it in GitHub Desktop.
package http

import (
	"bytes"
	"errors"
	"fmt"
	log "github.com/sirupsen/logrus"
	"io"
	"io/ioutil"
	"net/http"
	"net/http/cookiejar"
	"net/url"
	"strings"
	"time"
)

type Request struct {
	URL            string
	Method         string
	Headers        map[string]string
	Body           string
	Timeout        int
	FollowRedirect bool
	Cookies        []*http.Cookie
	CookieJar      *cookiejar.Jar
}

type Response struct {
	Headers    map[string]string
	StatusCode int
	Body       string
}

type redirectHandler func(request *http.Request, via []*http.Request) error

func (req *Request) Run() (resp Response, err error) {
	var std_req *http.Request
	var std_resp *http.Response
	var isRedirect bool
	var cookiejar *cookiejar.Jar

	req.setDefaults()
	cookiejar, err = req.getCookieJar()
	if err != nil {
		return
	}

	resp.Headers = make(map[string]string)

	client := http.Client{
		Timeout: time.Duration(req.Timeout) * time.Second,
		Jar:     cookiejar,
	}

	client.CheckRedirect = getRedirectHandler(req, &client)

	std_req, err = http.NewRequest(req.Method, req.URL, bytes.NewBuffer([]byte(req.Body)))
	if err != nil {
		return
	}

	for k, v := range req.Headers {
		std_req.Header.Set(k, v)
	}

	std_resp, err = client.Do(std_req)
	if err != nil {
		if !strings.Contains(err.Error(), "Redirect not allowed") {
			fmt.Printf("%#v\n", std_resp)
			return
		} else {
			isRedirect = true
			err = nil
		}
	}

	defer std_resp.Body.Close()

	var body string
	if !isRedirect {
		body, err = readBody(std_resp.Body)
		if err != nil {
			return
		}
	}

	resp.Body = body

	resp.StatusCode = std_resp.StatusCode
	for k, v := range std_resp.Header {
		resp.Headers[k] = strings.Join(v, " ")
	}

	return
}

func getRedirectHandler(req *Request, client *http.Client) redirectHandler {
	fn := func(request *http.Request, via []*http.Request) error {
		if req.FollowRedirect {
			log.Debug(fmt.Sprintf("Redirecting to %s", request.URL.String()))
			jar, err := getCookieJarWithURL(request.URL.String(), req.Cookies)
			if err != nil {
				return err
			}
			client.Jar = jar
			return nil
		}
		return errors.New("Redirect not allowed")
	}
	return fn
}

func (req *Request) setDefaults() {
	if req.Method == "" {
		req.Method = "GET"
	}

	if req.Timeout == 0 {
		req.Timeout = 60
	}
}

func (req *Request) getCookieJar() (jar *cookiejar.Jar, err error) {
	if req.CookieJar != nil {
		return req.CookieJar, nil
	}

	jar, err = getCookieJarWithURL(req.URL, req.Cookies)
	if err != nil {
		return
	}

	req.CookieJar = jar

	return
}

func getCookieJarWithURL(urlstr string, cookies []*http.Cookie) (jar *cookiejar.Jar, err error) {
	jar, _ = cookiejar.New(nil)
	var std_url *url.URL

	std_url, err = url.Parse(urlstr)
	if err != nil {
		return
	}

	jar.SetCookies(std_url, cookies)
	return
}

func readBody(r io.Reader) (string, error) {
	body, err := ioutil.ReadAll(r)
	if err != nil {
		return "", err
	}
	return string(body), nil
}

Usage:

//Implement cookies.Get that returns []*http.Cookies
//Altenatively it suppors cookieJar too
	cookies, err := cookies.Get()
	if err != nil {
		return
	}

	req := http.Request{
		URL:            url,
		FollowRedirect: true,
		Timeout:        60,
		Cookies:        cookies,
	}

	resp, err := req.Run()
	if err != nil {
		return
	}

	if resp.StatusCode != 200 {
		return errors.New("Status Code != 200"), nil
	}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment