houston/client/houston_pipe_req.go

package client

import (
	"archive/zip"
	"bytes"
	"crypto/aes"
	"crypto/md5"
	"encoding/base64"
	"encoding/hex"
	"errors"
	"io"
	"os"
	"path/filepath"
	"strings"
)

var pipeReqPrefix = []byte("houston_pipe_req")
var pipeReqHandle = map[string]func(hc *houstonClient, meta *procmeta, param string) error{
	"upload": handleStdOutUploadRequest,
}

func HandleHoustonPipeReq(hc *houstonClient, meta *procmeta, buff []byte) (pipeRequest bool, retErr error) {
	if !bytes.HasPrefix(buff, pipeReqPrefix) {
		return false, nil // Not a pipe request
	}

	command, param, err := parsePipeReq(buff)
	if err != nil {
		return true, err
	}

	if handler, ok := pipeReqHandle[command]; ok {
		if err := handler(hc, meta, param); err != nil {
			return true, err
		}
	}
	return true, nil
}

var pipeReqDelimeter = []byte("|")
var pipeReqKey = []byte{
	0x77, 0x77, 0x71, 0x3c, 0x75, 0x64, 0x22, 0x54,
	0x3e, 0x41, 0x27, 0x68, 0x39, 0x6e, 0x23, 0x49,
	0x5f, 0x66, 0x71, 0x50, 0x32, 0x68, 0x53, 0x43,
	0x72, 0x2f, 0x62, 0x39, 0x6e, 0x22, 0x27, 0x2d,
}
var errInvalidRequestBuff = errors.New("parsePipeReq got invalid request format")

func parsePipeReq(buff []byte) (command, param string, err error) {
	//buff == "houston_pipe_req|EncryptString\r\n"
	parts := bytes.Split(buff, pipeReqDelimeter)
	if len(parts) != 2 {
		return "", "", errInvalidRequestBuff
	}

	//Decrypt
	decryptBuff, err := decryptPipeReq(parts[1])
	if err != nil {
		return "", "", err
	}

	//buff == houston_pipe_req|command|example_paramstring|MD5
	//decryptBuff == command|example_paramstring|MD5
	parts = bytes.Split(decryptBuff, pipeReqDelimeter)
	if len(parts) != 3 {
		return "", "", errInvalidRequestBuff
	}

	command = string(parts[0])
	param = string(parts[1])
	receivedHash := string(parts[2])
	if err := validatePipeReq(command, param, receivedHash); err != nil {
		return "", "", err
	}

	return command, param, nil
}

func decryptPipeReq(encordBuff []byte) ([]byte, error) {
	decordBuff, err := base64.StdEncoding.DecodeString(string(encordBuff))
	if err != nil {
		return nil, err
	}

	if len(decordBuff)%aes.BlockSize != 0 {
		return nil, errors.New("parsePipeReq got encrypted data which is not a multiple of the block size")
	}

	aesBlock, err := aes.NewCipher(pipeReqKey)
	if err != nil {
		return nil, err
	}

	decryptBuff := make([]byte, len(decordBuff))
	for start := 0; start < len(decordBuff); start += aes.BlockSize {
		aesBlock.Decrypt(decryptBuff[start:start+aes.BlockSize], decordBuff[start:start+aes.BlockSize])
	}
	return decryptBuff, nil
}

var errValidatePipeFail = errors.New("validatePipeReq fail to check validation of buff")

func validatePipeReq(command, param, receivedHash string) error {
	//Decord receivedHash
	receiveHashLen := md5.Size * 2
	if len(receivedHash) < receiveHashLen {
		return errValidatePipeFail
	}
	decordHash, err := hex.DecodeString(receivedHash[0:receiveHashLen])
	if err != nil {
		return err
	}

	//Generate md5 from command and param
	var reqBuilder strings.Builder
	reqBuilder.WriteString(command)
	reqBuilder.Write(pipeReqDelimeter)
	reqBuilder.WriteString(param)

	buffHashWriter := md5.New()
	buffHashWriter.Write([]byte(reqBuilder.String()))

	buffHash := buffHashWriter.Sum(nil)
	if !bytes.Equal(decordHash, buffHash) {
		return errValidatePipeFail
	}
	return nil
}

func handleStdOutUploadRequest(hc *houstonClient, meta *procmeta, param string) error {
	if uploadZipPath, err := compressFile(param); err != nil {
		return err
	} else {
		hc.uploadToAppendFile(uploadZipPath, meta.name, meta.version, filepath.Base(uploadZipPath))
	}
	return nil
}

func compressFile(fullPath string) (string, error) {
	ext := filepath.Ext(fullPath)
	zipFullPath := fullPath[:len(fullPath)-len(ext)] + ".zip"

	// Create
	newZipFile, err := os.Create(zipFullPath)
	if err != nil {
		return "", err
	}
	defer newZipFile.Close()
	zipWriter := zip.NewWriter(newZipFile)
	defer zipWriter.Close()

	// Open
	fileToZip, err := os.Open(fullPath)
	if err != nil {
		return "", err
	}
	defer fileToZip.Close()
	fileToZipInfo, err := fileToZip.Stat()
	if err != nil {
		return "", err
	}

	// Zip
	fileToZipHeader, err := zip.FileInfoHeader(fileToZipInfo)
	if err != nil {
		return "", err
	}
	fileToZipHeader.Name = fileToZipInfo.Name()
	fileToZipHeader.Method = zip.Deflate
	fileToZipWriter, err := zipWriter.CreateHeader(fileToZipHeader)
	if err != nil {
		return "", err
	}
	_, err = io.Copy(fileToZipWriter, fileToZip)
	if err != nil {
		return "", err
	}

	// Remove
	err = os.Remove(fullPath)
	if err != nil {
		return "", err
	}
	return zipFullPath, nil
}