github.com/jfrog/jfrog-cli-go@v1.22.1-0.20200318093948-4826ef344ffd/bintray/helpers/streammanager.go (about)

     1  package helpers
     2  
     3  import (
     4  	"bufio"
     5  	"encoding/json"
     6  	"errors"
     7  	"github.com/jfrog/jfrog-cli-go/utils/cliutils"
     8  	"github.com/jfrog/jfrog-client-go/httpclient"
     9  	"github.com/jfrog/jfrog-client-go/utils/errorutils"
    10  	"github.com/jfrog/jfrog-client-go/utils/io/httputils"
    11  	"github.com/jfrog/jfrog-client-go/utils/log"
    12  	"io"
    13  	"io/ioutil"
    14  	"net/http"
    15  	"time"
    16  )
    17  
    18  const BINTRAY_RECONNECT_HEADER = "X-Bintray-Stream-Reconnect-Id"
    19  
    20  type StreamManager struct {
    21  	HttpClientDetails httputils.HttpClientDetails
    22  	Url               string
    23  	IncludeFilter     map[string]struct{}
    24  	ReconnectId       string
    25  }
    26  
    27  func (sm *StreamManager) ReadStream(resp *http.Response, writer io.Writer, lastServerInteraction *time.Time) {
    28  	ioReader := resp.Body
    29  	bodyReader := bufio.NewReader(ioReader)
    30  	sm.handleStream(bodyReader, writer, lastServerInteraction)
    31  }
    32  
    33  func (sm *StreamManager) handleStream(ioReader io.Reader, writer io.Writer, lastServerInteraction *time.Time) {
    34  	bodyReader := bufio.NewReader(ioReader)
    35  	pReader, pWriter := io.Pipe()
    36  	defer pWriter.Close()
    37  	go func() {
    38  		defer pReader.Close()
    39  		for {
    40  			line, _, err := bodyReader.ReadLine()
    41  			if err != nil {
    42  				log.Debug(err)
    43  				break
    44  			}
    45  			*lastServerInteraction = time.Now()
    46  			_, err = pWriter.Write(line)
    47  			if err != nil {
    48  				log.Debug(err)
    49  				break
    50  			}
    51  		}
    52  	}()
    53  	streamDecoder := json.NewDecoder(pReader)
    54  	streamEncoder := json.NewEncoder(writer)
    55  	sm.parseStream(streamDecoder, streamEncoder)
    56  }
    57  
    58  func (sm *StreamManager) parseStream(streamDecoder *json.Decoder, streamEncoder *json.Encoder) error {
    59  	for {
    60  		var decodedJson map[string]interface{}
    61  		if e := streamDecoder.Decode(&decodedJson); e != nil {
    62  			log.Debug(e)
    63  			return e
    64  		}
    65  		if _, ok := sm.IncludeFilter[decodedJson["type"].(string)]; ok || len(sm.IncludeFilter) == 0 {
    66  			if e := streamEncoder.Encode(&decodedJson); e != nil {
    67  				log.Debug(e)
    68  				return e
    69  			}
    70  		}
    71  	}
    72  }
    73  
    74  func (sm *StreamManager) isReconnection() bool {
    75  	return len(sm.ReconnectId) > 0
    76  }
    77  
    78  func (sm *StreamManager) setReconnectHeader() {
    79  	if sm.HttpClientDetails.Headers == nil {
    80  		sm.HttpClientDetails.Headers = map[string]string{}
    81  	}
    82  	sm.HttpClientDetails.Headers[BINTRAY_RECONNECT_HEADER] = sm.ReconnectId
    83  }
    84  
    85  func (sm *StreamManager) Connect() (bool, *http.Response) {
    86  	if sm.isReconnection() {
    87  		sm.setReconnectHeader()
    88  	}
    89  	log.Debug("Connecting...")
    90  	client, err := httpclient.ClientBuilder().Build()
    91  	if err != nil {
    92  		return false, nil
    93  	}
    94  	resp, _, _, e := client.Stream(sm.Url, sm.HttpClientDetails)
    95  	if e != nil {
    96  		return false, resp
    97  	}
    98  	if resp.StatusCode != http.StatusOK {
    99  		errorutils.CheckError(errors.New("response: " + resp.Status))
   100  		msgBody, _ := ioutil.ReadAll(resp.Body)
   101  		resp.Body.Close()
   102  		if resp.StatusCode > 400 && resp.StatusCode < 500 {
   103  			cliutils.ExitOnErr(errors.New(string(msgBody)))
   104  		}
   105  		return false, resp
   106  
   107  	}
   108  	sm.ReconnectId = resp.Header.Get(BINTRAY_RECONNECT_HEADER)
   109  	log.Debug("Connected.")
   110  	return true, resp
   111  }