waf-platform/EdgeDNS/internal/nodes/http_writer.go

// 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()
	}
}