/
PjaxFilter.go
104 lines (81 loc) · 2.22 KB
/
PjaxFilter.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
package pjax
import (
"fmt"
"io"
"net/http"
"code.google.com/p/go-html-transform/css/selector"
"code.google.com/p/go-html-transform/h5"
)
type PjaxFilter struct {
handler http.HandlerFunc
}
func NewPjaxFilter(handler http.HandlerFunc) *PjaxFilter {
return &PjaxFilter{
handler: handler,
}
}
func rewriteBody(containerSelector string, dest io.Writer, body string) (err error) {
if containerSelector == "" {
dest.Write([]byte(body))
return
}
var chain *selector.Chain
var document *h5.Tree
if document, err = h5.NewFromString(body); err != nil {
err = fmt.Errorf("invalid html document: %v", err)
return
}
var titleNode string
if titleNode, err = getTitleNode(document); err != nil {
return
}
if chain, err = selector.Selector(containerSelector); err != nil {
err = fmt.Errorf("invalid css: %v", containerSelector)
return
}
if matches := chain.Find(document.Top()); len(matches) > 0 {
match := matches[0:1] // Take only the first match
newBody := h5.RenderNodesToString(h5.Children(match[0]))
fmt.Printf("data: %v", h5.Data(match[0]))
dest.Write([]byte(titleNode))
dest.Write([]byte(newBody))
return
}
err = fmt.Errorf("container not found")
return
}
func getTitleNode(document *h5.Tree) (titleNode string, err error) {
var chain *selector.Chain
if chain, err = selector.Selector("title"); err != nil {
return
}
if matches := chain.Find(document.Top()); len(matches) > 0 {
match := matches[0:1]
titleNode = h5.RenderNodesToString(match)
}
return
}
func (s *PjaxFilter) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
var container string
// Validate pjax meta data.
if container = req.Header.Get("X-PJAX-CONTAINER"); container == "" {
container = req.URL.Query().Get("_pjax")
}
// We didn't find a container, so handle as non pjax request.
if container == "" {
// Request contains no pjax information.
s.handler(rw, req)
return
}
response := NewResponseRecorder()
s.handler(response, req)
for key, _ := range response.Header() {
if key != "Content-Lenght" {
value := response.Header().Get(key)
response.Header().Set(key, value)
}
}
if rewriteErr := rewriteBody(container, rw, response.body.String()); rewriteErr != nil {
response.body.WriteTo(rw)
}
}