linkinfo/default.go

package linkinfo

import (
	"errors"
	"fmt"
	"github.com/PuerkitoBio/goquery"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"path"
	"strconv"
	"strings"
)

const (
	contentTypeHtml = "text/html"
)

func defaultLinkHandler(link string) (*LinkInfo, error) {
	redirects := make([]string, 0)

	u, err := url.Parse(link)

	if err != nil {
		return nil, err
	}

	var res *http.Response

	for i := 0; i < 10; i++ {
		res, err = Client.Head(link)

		if err != nil {
			return nil, err
		}

		if (res.StatusCode == 301 || res.StatusCode == 302) && res.Header.Get("Location") != "" {
			link = res.Header.Get("Location")
			redirects = append(redirects, link)
		} else {
			break
		}
	}

	if res != nil && res.StatusCode != 200 {
		return nil, errors.New("invalid response, expected 200, got " + strconv.Itoa(res.StatusCode))
	}

	contentType := res.Header.Get("Content-Type")

	if contentType == "" {
		contentType = detectContentType(link, "application/octet-stream")
	}

	if idx := strings.Index(contentType, ";"); idx != -1 {
		contentType = contentType[:idx]
	}

	var contentLength int64

	if contentLengthStr := res.Header.Get("Content-Length"); contentLengthStr != "" {
		contentLength, err = strconv.ParseInt(contentLengthStr, 10, 64)
	}

	ret := &LinkInfo{
		ContentType:   contentType,
		ContentLength: contentLength,
	}

	switch contentType {
	case contentTypeHtml:
		err = retrieveHtmlLinkTitle(ret, link)
	default:
		ret.Title = fmt.Sprintf("%s (%s, %s)", path.Base(u.Path), contentType, ByteCountDecimal(contentLength))
	}

	return ret, err
}

func detectContentType(link, defaultType string) string {
	req, err := http.NewRequest("GET", link, nil)

	if err != nil {
		return defaultType
	}

	req.Header.Set("Range", "bytes=0-512")

	res, err := Client.Do(req)

	if err != nil {
		return defaultType
	}

	defer res.Body.Close()

	b, err := ioutil.ReadAll(io.LimitReader(res.Body, 512))

	if err != nil {
		return defaultType
	}

	t := http.DetectContentType(b)

	if t == "" {
		t = defaultType
	}

	return t
}

var (
	attrKeys = []string{"property", "name", "itemprop"}
)

func retrieveHtmlLinkTitle(i *LinkInfo, link string) error {
	res, err := Client.Get(link)

	if err != nil {
		return err
	}

	defer res.Body.Close()

	q, err := goquery.NewDocumentFromReader(res.Body)

	if err != nil {
		return err
	}

	meta := q.Find("meta")

	metaTags := make(map[string]string)

	meta.Each(func(_ int, s *goquery.Selection) {
		var key string
		var exists bool

		for _, k := range attrKeys {
			key, exists = s.Attr(k)

			if exists {
				break
			}
		}

		if key == "" {
			return
		}
	})

	var attr string
	var exists bool

	if attr, exists = metaTags["og:title"]; exists {
		i.Title = attr
	} else if tag := q.Find("title"); tag.Length() > 0 {
		i.Title = strings.TrimSpace(tag.Text())
	}

	if attr, exists = metaTags["og:description"]; exists {
		i.Description = attr
	} else if attr, exists = metaTags["description"]; exists {
		i.Description = attr
	}

	if attr, exists = metaTags["duration"]; exists {
		i.Duration = attr
	}

	return nil
}