fix download part problem, add png support

1 files changed, 162 insertions, 17 deletions

diff --git a/modules/avatar/avatar.go b/modules/avatar/avatar.go
index 93f842eaef..55d1e13d94 100644
--- a/modules/avatar/avatar.go
+++ b/modules/avatar/avatar.go
@@ -3,7 +3,11 @@ package avatar
 import (
 	"crypto/md5"
 	"encoding/hex"
+	"errors"
 	"fmt"
+	"image"
+	"image/jpeg"
+	"image/png"
 	"io"
 	"log"
 	"net/http"
@@ -13,9 +17,14 @@ import (
 	"strings"
 	"sync"
 	"time"
+
+	"github.com/nfnt/resize"
 )
 
-var gravatar = "http://www.gravatar.com/avatar"
+var (
+	gravatar         = "http://www.gravatar.com/avatar"
+	defaultImagePath = "./default.jpg"
+)
 
 // hash email to md5 string
 func HashEmail(email string) string {
@@ -25,37 +34,145 @@ func HashEmail(email string) string {
 }
 
 type Avatar struct {
-	Hash      string
-	cacheDir  string // image save dir
-	reqParams string
-	imagePath string
+	Hash           string
+	cacheDir       string // image save dir
+	reqParams      string
+	imagePath      string
+	expireDuration time.Duration
 }
 
 func New(hash string, cacheDir string) *Avatar {
 	return &Avatar{
-		Hash:     hash,
-		cacheDir: cacheDir,
+		Hash:           hash,
+		cacheDir:       cacheDir,
+		expireDuration: time.Minute * 10,
 		reqParams: url.Values{
 			"d":    {"retro"},
 			"size": {"200"},
 			"r":    {"pg"}}.Encode(),
-		imagePath: filepath.Join(cacheDir, hash+".jpg"),
+		imagePath: filepath.Join(cacheDir, hash+".image"), //maybe png or jpeg
+	}
+}
+
+func (this *Avatar) InCache() bool {
+	fileInfo, err := os.Stat(this.imagePath)
+	return err == nil && fileInfo.Mode().IsRegular()
+}
+
+func (this *Avatar) Modtime() (modtime time.Time, err error) {
+	fileInfo, err := os.Stat(this.imagePath)
+	if err != nil {
+		return
+	}
+	return fileInfo.ModTime(), nil
+}
+
+func (this *Avatar) Expired() bool {
+	if !this.InCache() {
+		return true
+	}
+	fileInfo, err := os.Stat(this.imagePath)
+	return err != nil || time.Since(fileInfo.ModTime()) > this.expireDuration
+}
+
+// default image format: jpeg
+func (this *Avatar) Encode(wr io.Writer, size int) (err error) {
+	var img image.Image
+	decodeImageFile := func(file string) (img image.Image, err error) {
+		fd, err := os.Open(file)
+		if err != nil {
+			return
+		}
+		defer fd.Close()
+		img, err = jpeg.Decode(fd)
+		if err != nil {
+			fd.Seek(0, os.SEEK_SET)
+			img, err = png.Decode(fd)
+		}
+		return
+	}
+	imgPath := this.imagePath
+	if !this.InCache() {
+		imgPath = defaultImagePath
+	}
+	img, err = decodeImageFile(imgPath)
+	if err != nil {
+		return
 	}
+	m := resize.Resize(uint(size), 0, img, resize.Lanczos3)
+	return jpeg.Encode(wr, m, nil)
 }
 
 // get image from gravatar.com
 func (this *Avatar) Update() {
 	thunder.Fetch(gravatar+"/"+this.Hash+"?"+this.reqParams,
-		this.Hash+".jpg")
+		this.imagePath)
 }
 
-func (this *Avatar) UpdateTimeout(timeout time.Duration) {
+func (this *Avatar) UpdateTimeout(timeout time.Duration) error {
+	var err error
 	select {
 	case <-time.After(timeout):
-		log.Println("timeout")
-	case <-thunder.GoFetch(gravatar+"/"+this.Hash+"?"+this.reqParams,
-		this.Hash+".jpg"):
+		err = errors.New("get gravatar image timeout")
+	case err = <-thunder.GoFetch(gravatar+"/"+this.Hash+"?"+this.reqParams,
+		this.imagePath):
 	}
+	return err
+}
+
+func init() {
+	log.SetFlags(log.Lshortfile | log.LstdFlags)
+}
+
+// http.Handle("/avatar/", avatar.HttpHandler("./cache"))
+func HttpHandler(cacheDir string) func(w http.ResponseWriter, r *http.Request) {
+	MustInt := func(r *http.Request, defaultValue int, keys ...string) int {
+		var v int
+		for _, k := range keys {
+			if _, err := fmt.Sscanf(r.FormValue(k), "%d", &v); err == nil {
+				defaultValue = v
+			}
+		}
+		return defaultValue
+	}
+
+	return func(w http.ResponseWriter, r *http.Request) {
+		urlPath := r.URL.Path
+		hash := urlPath[strings.LastIndex(urlPath, "/")+1:]
+		hash = HashEmail(hash)
+		size := MustInt(r, 80, "s", "size") // size = 80*80
+
+		avatar := New(hash, cacheDir)
+		if avatar.Expired() {
+			err := avatar.UpdateTimeout(time.Millisecond * 500)
+			if err != nil {
+				log.Println(err)
+			}
+		}
+		if modtime, err := avatar.Modtime(); err == nil {
+			etag := fmt.Sprintf("size(%d)", size)
+			if t, err := time.Parse(http.TimeFormat, r.Header.Get("If-Modified-Since")); err == nil && modtime.Before(t.Add(1*time.Second)) && etag == r.Header.Get("If-None-Match") {
+				h := w.Header()
+				delete(h, "Content-Type")
+				delete(h, "Content-Length")
+				w.WriteHeader(http.StatusNotModified)
+				return
+			}
+			w.Header().Set("Last-Modified", modtime.UTC().Format(http.TimeFormat))
+			w.Header().Set("ETag", etag)
+		}
+		w.Header().Set("Content-Type", "image/jpeg")
+		err := avatar.Encode(w, size)
+		if err != nil {
+			log.Println(err)
+			w.WriteHeader(500)
+		}
+	}
+}
+
+func init() {
+	http.HandleFunc("/", HttpHandler("./"))
+	log.Fatal(http.ListenAndServe(":8001", nil))
 }
 
 var thunder = &Thunder{QueueSize: 10}
@@ -114,8 +231,17 @@ func (this *thunderTask) Fetch() {
 	this.Done()
 }
 
+var client = &http.Client{}
+
 func (this *thunderTask) fetch() error {
-	resp, err := http.Get(this.Url)
+	log.Println("thunder, fetch", this.Url)
+	req, _ := http.NewRequest("GET", this.Url, nil)
+	req.Header.Set("Accept", "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8")
+	req.Header.Set("Accept-Encoding", "gzip,deflate,sdch")
+	req.Header.Set("Accept-Language", "zh-CN,zh;q=0.8")
+	req.Header.Set("Cache-Control", "no-cache")
+	req.Header.Set("User-Agent", "Mozilla/5.0 (Windows NT 6.1; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/33.0.1750.154 Safari/537.36")
+	resp, err := client.Do(req)
 	if err != nil {
 		return err
 	}
@@ -123,14 +249,33 @@ func (this *thunderTask) fetch() error {
 	if resp.StatusCode != 200 {
 		return fmt.Errorf("status code: %d", resp.StatusCode)
 	}
-	fd, err := os.Create(this.SaveFile)
+
+	/*
+		log.Println("headers:", resp.Header)
+		switch resp.Header.Get("Content-Type") {
+		case "image/jpeg":
+			this.SaveFile += ".jpeg"
+		case "image/png":
+			this.SaveFile += ".png"
+		}
+	*/
+	/*
+		imgType := resp.Header.Get("Content-Type")
+		if imgType != "image/jpeg" && imgType != "image/png" {
+			return errors.New("not png or jpeg")
+		}
+	*/
+
+	tmpFile := this.SaveFile + ".part" // mv to destination when finished
+	fd, err := os.Create(tmpFile)
 	if err != nil {
 		return err
 	}
-	defer fd.Close()
 	_, err = io.Copy(fd, resp.Body)
+	fd.Close()
 	if err != nil {
+		os.Remove(tmpFile)
 		return err
 	}
-	return nil
+	return os.Rename(tmpFile, this.SaveFile)
 }