Changeset 45 in code

Timestamp:

Nov 28, 2016, 12:49:22 PM (9 years ago)

Author:

asciimoo

Message:

Merge pull request #35 from dalf/url_encoding

Fetched URI matched the mortyurl.

File:

• trunk/morty.go (modified) (10 diffs)

trunk/morty.go

@@ -18 +18 @@
         "github.com/valyala/fasthttp"
         "golang.org/x/net/html"
-        "golang.org/x/text/encoding/charmap"
+        "golang.org/x/net/html/charset"
+        "golang.org/x/text/encoding"
 )
 
@@ -123 +124 @@
 `
 
+var HTML_META_CONTENT_TYPE string = "<meta http-equiv=\"Content-Type\" content=\"text/html; charset=utf-8\">"
+
 func (p *Proxy) RequestHandler(ctx *fasthttp.RequestCtx) {
 
@@ -230 +233 @@
         var responseBody []byte
 
-        if len(contentInfo) == 2 && bytes.Contains(contentInfo[1], []byte("ISO-8859-2")) && bytes.Contains(contentInfo[0], []byte("text")) {
-                var err error
-                responseBody, err = charmap.ISO8859_2.NewDecoder().Bytes(resp.Body())
-                if err != nil {
-                        // HTTP status code 503 : Service Unavailable
-                        p.serveMainPage(ctx, 503, err)
-                        return
+        if len(contentInfo) == 2 && bytes.Contains(contentInfo[0], []byte("text")) {
+                e, ename, _ := charset.DetermineEncoding(resp.Body(), string(contentType))
+                if (e != encoding.Nop) && (!strings.EqualFold("utf-8", ename)) {
+                        responseBody, err = e.NewDecoder().Bytes(resp.Body())
+                        if err != nil {
+                                // HTTP status code 503 : Service Unavailable
+                                p.serveMainPage(ctx, 503, err)
+                                return
+                        }
+                } else {
+                        responseBody = resp.Body()
                 }
         } else {
@@ -319 +326 @@
         unsafeElements := make([][]byte, 0, 8)
         state := STATE_DEFAULT
-
         for {
                 token := decoder.Next()
@@ -347 +353 @@
                                         for {
                                                 attrName, attrValue, moreAttr := decoder.TagAttr()
-                                                if !bytes.Equal(attrName, []byte("href")) {
-                                                        continue
-                                                }
-                                                parsedURI, err := url.Parse(string(attrValue))
-                                                if err == nil {
-                                                        rc.BaseURL = parsedURI
+                                                if bytes.Equal(attrName, []byte("href")) {
+                                                        parsedURI, err := url.Parse(string(attrValue))
+                                                        if err == nil {
+                                                                rc.BaseURL = parsedURI
+                                                        }
                                                 }
                                                 if !moreAttr {
@@ -383 +388 @@
                                 }
 
+                                if bytes.Equal(tag, []byte("meta")) {
+                                        sanitizeMetaTag(rc, out, attrs)
+                                        break
+                                }
+
                                 fmt.Fprintf(out, "<%s", tag)
 
                                 if hasAttrs {
-                                        if bytes.Equal(tag, []byte("meta")) {
-                                                sanitizeMetaAttrs(rc, out, attrs)
-                                        } else {
-                                                sanitizeAttrs(rc, out, attrs)
-                                        }
+                                        sanitizeAttrs(rc, out, attrs)
                                 }
 
@@ -400 +406 @@
                                                 state = STATE_IN_STYLE
                                         }
+                                }
+
+                                if bytes.Equal(tag, []byte("head")) {
+                                        fmt.Fprintf(out, HTML_META_CONTENT_TYPE)
                                 }
 
@@ -499 +509 @@
 }
 
-func sanitizeMetaAttrs(rc *RequestConfig, out io.Writer, attrs [][][]byte) {
+func sanitizeMetaTag(rc *RequestConfig, out io.Writer, attrs [][][]byte) {
         var http_equiv []byte
         var content []byte
@@ -512 +522 @@
                         content = attrValue
                 }
-        }
-
+                if bytes.Equal(attrName, []byte("charset")) {
+                        // exclude <meta charset="...">
+                        return
+                }
+        }
+
+        if bytes.Equal(http_equiv, []byte("content-type")) {
+                return
+        }
+
+        out.Write([]byte("<meta"))
         urlIndex := bytes.Index(bytes.ToLower(content), []byte("url="))
         if bytes.Equal(http_equiv, []byte("refresh")) && urlIndex != -1 {
@@ -530 +549 @@
                 sanitizeAttrs(rc, out, attrs)
         }
-
+        out.Write([]byte(">"))
 }