1.4.5.2

EdgeDNS/internal/nodes/http_writer.go

@@ -0,0 +1,179 @@
+// Copyright 2023 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn .
+
+package nodes
+
+import (
+	"encoding/json"
+	"errors"
+	"github.com/iwind/TeaGo/Tea"
+	"github.com/iwind/TeaGo/types"
+	"github.com/miekg/dns"
+	"net"
+	"net/http"
+	"reflect"
+)
+
+type HTTPWriter struct {
+	rawConn     net.Conn
+	rawWriter   http.ResponseWriter
+	contentType string
+}
+
+func NewHTTPWriter(rawWriter http.ResponseWriter, rawConn net.Conn, contentType string) *HTTPWriter {
+	return &HTTPWriter{
+		rawWriter:   rawWriter,
+		rawConn:     rawConn,
+		contentType: contentType,
+	}
+}
+
+func (this *HTTPWriter) LocalAddr() net.Addr {
+	return this.rawConn.LocalAddr()
+}
+
+func (this *HTTPWriter) RemoteAddr() net.Addr {
+	return this.rawConn.RemoteAddr()
+}
+
+func (this *HTTPWriter) WriteMsg(msg *dns.Msg) error {
+	if msg == nil {
+		return errors.New("'msg' should not be nil")
+	}
+
+	msgData, err := this.encodeMsg(msg)
+	if err != nil {
+		return err
+	}
+
+	this.rawWriter.Header().Set("Content-Length", types.String(len(msgData)))
+	this.rawWriter.Header().Set("Content-Type", this.contentType)
+
+	// cache-control
+	if len(msg.Answer) > 0 {
+		var minTtl uint32
+		for _, answer := range msg.Answer {
+			var header = answer.Header()
+			if header != nil && header.Ttl > 0 && (minTtl == 0 || header.Ttl < minTtl) {
+				minTtl = header.Ttl
+			}
+		}
+		if minTtl > 0 {
+			this.rawWriter.Header().Set("Cache-Control", "max-age="+types.String(minTtl))
+		}
+	}
+
+	this.rawWriter.WriteHeader(http.StatusOK)
+	_, err = this.rawWriter.Write(msgData)
+	return err
+}
+
+func (this *HTTPWriter) Write(p []byte) (int, error) {
+	this.rawWriter.Header().Set("Content-Length", types.String(len(p)))
+	this.rawWriter.WriteHeader(http.StatusOK)
+	return this.rawWriter.Write(p)
+}
+
+func (this *HTTPWriter) Close() error {
+	return nil
+}
+
+func (this *HTTPWriter) TsigStatus() error {
+	return nil
+}
+
+func (this *HTTPWriter) TsigTimersOnly(timersOnly bool) {
+}
+
+func (this *HTTPWriter) Hijack() {
+	hijacker, ok := this.rawWriter.(http.Hijacker)
+	if ok {
+		_, _, _ = hijacker.Hijack()
+	}
+}
+
+func (this *HTTPWriter) encodeMsg(msg *dns.Msg) ([]byte, error) {
+	if this.contentType == "application/x-javascript" || this.contentType == "application/json" {
+		var result = map[string]any{
+			"Status": 0,
+			"TC":     msg.Truncated,
+			"RD":     msg.RecursionDesired,
+			"RA":     msg.RecursionAvailable,
+			"AD":     msg.AuthenticatedData,
+			"CD":     msg.CheckingDisabled,
+		}
+
+		// questions
+		var questionMaps = []map[string]any{}
+		for _, question := range msg.Question {
+			questionMaps = append(questionMaps, map[string]any{
+				"name": question.Name,
+				"type": question.Qtype,
+			})
+		}
+		result["Question"] = questionMaps
+
+		// answers
+		var answerMaps = []map[string]any{}
+		for _, answer := range msg.Answer {
+			var answerMap = map[string]any{
+				"name": answer.Header().Name,
+				"type": answer.Header().Rrtype,
+				"TTL":  answer.Header().Ttl,
+			}
+
+			switch x := answer.(type) {
+			case *dns.A:
+				answerMap["data"] = x.A.String()
+			case *dns.AAAA:
+				answerMap["data"] = x.AAAA.String()
+			case *dns.CNAME:
+				answerMap["data"] = x.Target
+			case *dns.TXT:
+				answerMap["data"] = x.Txt
+			case *dns.NS:
+				answerMap["data"] = x.Ns
+			case *dns.MX:
+				answerMap["data"] = x.Mx
+				answerMap["preference"] = x.Preference
+			default:
+				var answerValue = reflect.ValueOf(answer).Elem()
+				var answerType = answerValue.Type()
+
+				var countFields = answerType.NumField()
+				for i := 0; i < countFields; i++ {
+					var fieldName = answerType.Field(i).Name
+					var fieldValue = answerValue.FieldByName(fieldName)
+					if !fieldValue.IsValid() {
+						continue
+					}
+					var fieldInterface = fieldValue.Interface()
+					if fieldInterface == nil {
+						continue
+					}
+
+					_, ok := fieldInterface.(dns.RR_Header)
+					if ok {
+						continue
+					}
+
+					if countFields == 2 {
+						answerMap["data"] = fieldValue.Interface()
+					} else {
+						answerMap[fieldName] = fieldValue.Interface()
+					}
+				}
+			}
+
+			answerMaps = append(answerMaps, answerMap)
+		}
+		result["Answer"] = answerMaps
+
+		if Tea.IsTesting() {
+			return json.MarshalIndent(result, "", "  ")
+		} else {
+			return json.Marshal(result)
+		}
+	} else {
+		return msg.Pack()
+	}
+}