Use context for more functions

This commit is contained in:
Star Brilliant committed 2018-11-07 19:25:46 +08:00
1 parent 95ec839409
commit c51be0e69c
6 files changed
+28 -24

No files matched your search

+7 -4
View File
@@ -217,6 +217,9 @@ func (c *Client) Start() error {
}
func (c *Client) handlerFunc(w dns.ResponseWriter, r *dns.Msg, isTCP bool) {
ctx, cancel := context.WithTimeout(context.Background(), time.Duration(c.conf.Timeout)*time.Second)
defer cancel()
if r.Response == true {
log.Println("Received a response packet")
return
@@ -298,9 +301,9 @@ func (c *Client) handlerFunc(w dns.ResponseWriter, r *dns.Msg, isTCP bool) {
var req *DNSRequest
if requestType == "application/dns-json" {
req = c.generateRequestGoogle(w, r, isTCP)
req = c.generateRequestGoogle(ctx, w, r, isTCP)
} else if requestType == "application/dns-message" {
req = c.generateRequestIETF(w, r, isTCP)
req = c.generateRequestIETF(ctx, w, r, isTCP)
} else {
panic("Unknown request Content-Type")
}
@@ -334,9 +337,9 @@ func (c *Client) handlerFunc(w dns.ResponseWriter, r *dns.Msg, isTCP bool) {
}
if contentType == "application/json" {
c.parseResponseGoogle(w, r, isTCP, req)
c.parseResponseGoogle(ctx, w, r, isTCP, req)
} else if contentType == "application/dns-message" {
c.parseResponseIETF(w, r, isTCP, req)
c.parseResponseIETF(ctx, w, r, isTCP, req)
} else {
panic("Unknown response Content-Type")
}
+2 -5
View File
@@ -34,13 +34,12 @@ import (
"net/url"
"strconv"
"strings"
"time"
"github.com/m13253/dns-over-https/json-dns"
"github.com/miekg/dns"
)
func (c *Client) generateRequestGoogle(w dns.ResponseWriter, r *dns.Msg, isTCP bool) *DNSRequest {
func (c *Client) generateRequestGoogle(ctx context.Context, w dns.ResponseWriter, r *dns.Msg, isTCP bool) *DNSRequest {
question := &r.Question[0]
questionName := question.Name
questionClass := question.Qclass
@@ -89,12 +88,10 @@ func (c *Client) generateRequestGoogle(w dns.ResponseWriter, r *dns.Msg, isTCP b
}
req.Header.Set("Accept", "application/json, application/dns-message, application/dns-udpwireformat")
req.Header.Set("User-Agent", USER_AGENT)
ctx, cancel := context.WithTimeout(context.Background(), time.Duration(c.conf.Timeout)*time.Second)
req = req.WithContext(ctx)
c.httpClientMux.RLock()
resp, err := c.httpClient.Do(req)
c.httpClientMux.RUnlock()
cancel()
if err != nil {
log.Println(err)
reply := jsonDNS.PrepareReply(r)
@@ -115,7 +112,7 @@ func (c *Client) generateRequestGoogle(w dns.ResponseWriter, r *dns.Msg, isTCP b
}
}
func (c *Client) parseResponseGoogle(w dns.ResponseWriter, r *dns.Msg, isTCP bool, req *DNSRequest) {
func (c *Client) parseResponseGoogle(ctx context.Context, w dns.ResponseWriter, r *dns.Msg, isTCP bool, req *DNSRequest) {
if req.response.StatusCode != 200 {
log.Printf("HTTP error from upstream %s: %s\n", req.currentUpstream, req.response.Status)
req.reply.Rcode = dns.RcodeServerFailure
+2 -4
View File
@@ -40,7 +40,7 @@ import (
"github.com/miekg/dns"
)
func (c *Client) generateRequestIETF(w dns.ResponseWriter, r *dns.Msg, isTCP bool) *DNSRequest {
func (c *Client) generateRequestIETF(ctx context.Context, w dns.ResponseWriter, r *dns.Msg, isTCP bool) *DNSRequest {
opt := r.IsEdns0()
udpSize := uint16(512)
if opt == nil {
@@ -127,12 +127,10 @@ func (c *Client) generateRequestIETF(w dns.ResponseWriter, r *dns.Msg, isTCP boo
}
req.Header.Set("Accept", "application/dns-message, application/dns-udpwireformat, application/json")
req.Header.Set("User-Agent", USER_AGENT)
ctx, cancel := context.WithTimeout(context.Background(), time.Duration(c.conf.Timeout)*time.Second)
req = req.WithContext(ctx)
c.httpClientMux.RLock()
resp, err := c.httpClient.Do(req)
c.httpClientMux.RUnlock()
cancel()
if err != nil {
log.Println(err)
reply := jsonDNS.PrepareReply(r)
@@ -153,7 +151,7 @@ func (c *Client) generateRequestIETF(w dns.ResponseWriter, r *dns.Msg, isTCP boo
}
}
func (c *Client) parseResponseIETF(w dns.ResponseWriter, r *dns.Msg, isTCP bool, req *DNSRequest) {
func (c *Client) parseResponseIETF(ctx context.Context, w dns.ResponseWriter, r *dns.Msg, isTCP bool, req *DNSRequest) {
if req.response.StatusCode != 200 {
log.Printf("HTTP error from upstream %s: %s\n", req.currentUpstream, req.response.Status)
req.reply.Rcode = dns.RcodeServerFailure