github.com/dbernstein1/tyk@v2.9.0-beta9-dl-apic+incompatible/gateway/res_handler_transform.go (about)

     1  package gateway
     2  
     3  import (
     4  	"bytes"
     5  	"compress/flate"
     6  	"compress/gzip"
     7  	"encoding/json"
     8  	"io"
     9  	"io/ioutil"
    10  	"net/http"
    11  	"strconv"
    12  
    13  	"github.com/sirupsen/logrus"
    14  	"github.com/clbanning/mxj"
    15  
    16  	"github.com/TykTechnologies/tyk/apidef"
    17  	"github.com/TykTechnologies/tyk/headers"
    18  	"github.com/TykTechnologies/tyk/user"
    19  )
    20  
    21  type ResponseTransformMiddleware struct {
    22  	Spec *APISpec
    23  }
    24  
    25  func (ResponseTransformMiddleware) Name() string {
    26  	return "ResponseTransformMiddleware"
    27  }
    28  
    29  func (h *ResponseTransformMiddleware) Init(c interface{}, spec *APISpec) error {
    30  	h.Spec = spec
    31  	return nil
    32  }
    33  
    34  func respBodyReader(req *http.Request, resp *http.Response) io.ReadCloser {
    35  
    36  	if req.Header.Get(headers.AcceptEncoding) == "" {
    37  		return resp.Body
    38  	}
    39  
    40  	switch resp.Header.Get(headers.ContentEncoding) {
    41  	case "gzip":
    42  		reader, err := gzip.NewReader(resp.Body)
    43  		if err != nil {
    44  			log.Error("Body decompression error:", err)
    45  			return ioutil.NopCloser(bytes.NewReader(nil))
    46  		}
    47  
    48  		// represents unknown length
    49  		resp.ContentLength = 0
    50  
    51  		return reader
    52  	case "deflate":
    53  		// represents unknown length
    54  		resp.ContentLength = 0
    55  
    56  		return flate.NewReader(resp.Body)
    57  	}
    58  
    59  	return resp.Body
    60  }
    61  
    62  func compressBuffer(in bytes.Buffer, encoding string) (out bytes.Buffer) {
    63  	switch encoding {
    64  	case "gzip":
    65  		zw := gzip.NewWriter(&out)
    66  		zw.Write(in.Bytes())
    67  		zw.Close()
    68  	case "deflate":
    69  		zw, _ := flate.NewWriter(&out, 1)
    70  		zw.Write(in.Bytes())
    71  		zw.Close()
    72  	default:
    73  		out = in
    74  	}
    75  
    76  	return out
    77  }
    78  
    79  func (h *ResponseTransformMiddleware) HandleResponse(rw http.ResponseWriter, res *http.Response, req *http.Request, ses *user.SessionState) error {
    80  
    81  	logger := log.WithFields(logrus.Fields{
    82  		"prefix":      "outbound-transform",
    83  		"server_name": h.Spec.Proxy.TargetURL,
    84  		"api_id":      h.Spec.APIID,
    85  		"path":        req.URL.Path,
    86  	})
    87  
    88  	_, versionPaths, _, _ := h.Spec.Version(req)
    89  	found, meta := h.Spec.CheckSpecMatchesStatus(req, versionPaths, TransformedResponse)
    90  	if !found {
    91  		return nil
    92  	}
    93  	tmeta := meta.(*TransformSpec)
    94  
    95  	respBody := respBodyReader(req, res)
    96  	body, _ := ioutil.ReadAll(respBody)
    97  	defer respBody.Close()
    98  
    99  	// Put into an interface:
   100  	bodyData := make(map[string]interface{})
   101  	switch tmeta.TemplateData.Input {
   102  	case apidef.RequestXML:
   103  		if len(body) == 0 {
   104  			body = []byte("<_/>")
   105  		}
   106  
   107  		mxj.XmlCharsetReader = WrappedCharsetReader
   108  		var err error
   109  
   110  		xmlMap, err := mxj.NewMapXml(body) // unmarshal
   111  		if err != nil {
   112  			logger.WithError(err).Error("Error unmarshalling XML")
   113  			//todo return error
   114  			break
   115  		}
   116  		for k, v := range xmlMap {
   117  			bodyData[k] = v
   118  		}
   119  	default: // apidef.RequestJSON
   120  		if len(body) == 0 {
   121  			body = []byte("{}")
   122  		}
   123  
   124  		var tempBody interface{}
   125  		if err := json.Unmarshal(body, &tempBody); err != nil {
   126  			logger.WithError(err).Error("Error unmarshalling JSON")
   127  			//todo return error
   128  			break
   129  		}
   130  
   131  		switch tempBody.(type) {
   132  		case []interface{}:
   133  			bodyData["array"] = tempBody
   134  		case map[string]interface{}:
   135  			bodyData = tempBody.(map[string]interface{})
   136  		}
   137  	}
   138  
   139  	if h.Spec.EnableContextVars {
   140  		bodyData["_tyk_context"] = ctxGetData(req)
   141  	}
   142  
   143  	if tmeta.TemplateData.EnableSession {
   144  		if session := ctxGetSession(req); session != nil {
   145  			bodyData["_tyk_meta"] = session.MetaData
   146  		} else {
   147  			logger.Error("Session context was enabled but not found.")
   148  		}
   149  	}
   150  
   151  	// Apply to template
   152  	var bodyBuffer bytes.Buffer
   153  	if err := tmeta.Template.Execute(&bodyBuffer, bodyData); err != nil {
   154  		logger.WithError(err).Error("Failed to apply template to request")
   155  	}
   156  
   157  	// Re-compress if original upstream response was compressed
   158  	encoding := res.Header.Get("Content-Encoding")
   159  	bodyBuffer = compressBuffer(bodyBuffer, encoding)
   160  
   161  	res.ContentLength = int64(bodyBuffer.Len())
   162  	res.Header.Set("Content-Length", strconv.Itoa(bodyBuffer.Len()))
   163  	res.Body = ioutil.NopCloser(&bodyBuffer)
   164  
   165  	return nil
   166  }