package cnp import ( "bufio" "bytes" "io" "io/ioutil" "strings" ) // Message represents a CNP message. type Message struct { Header Header Body io.Reader closer io.Closer } // NewMessage creates a new CNP message. // // This method also calls Message.TryComputeLength(). func NewMessage(intent string, body io.Reader) *Message { msg := &Message{ Header: NewHeader(intent, nil), Body: body, } msg.TryComputeLength() return msg } // ParseMessage parses a CNP message. // // The message's Body field is set to a bufio.Reader wrapping r. If r is an // io.Closer, it is also stored separately for usage with Message.Close(). func ParseMessage(r io.Reader) (*Message, error) { br := bufio.NewReader(r) line, err := readLimitedLine(br, MaxHeaderLength) if err != nil { return nil, err } h, err := ParseHeader(line) if err != nil { return nil, err } msg := &Message{ Header: h, Body: br, } if rc, ok := r.(io.Closer); ok { msg.closer = rc } return msg, nil } // Close attempts to close the message body reader. // // If the message's Body field is an io.Closer, its Close() method is called. // If Body is not an io.Closer and the message was created with ParseMessage // provided with an io.Closer, the Close() method on the original reader will // be called. // Otherwise, this function does nothing. func (msg *Message) Close() error { if rc, ok := msg.Body.(io.Closer); ok { return rc.Close() } if msg.closer != nil { return msg.closer.Close() } return nil } // ComputeLength sets the length header parameter based on the message body. // First, msg.TryComputeLength() is attempted; if that fails, the request is // fully read into a bytes.Reader and msg.Body is set to it. func (msg *Message) ComputeLength() error { if !msg.TryComputeLength() { buf, err := ioutil.ReadAll(msg.Body) if len(buf) > 0 { msg.Body = bytes.NewReader(buf) msg.SetLength(int64(len(buf))) } if err != nil { return err } } return nil } // TryComputeLength sets the length header parameter to the length of the // message body if the body's type is one of *bytes.Buffer, *bytes.Reader or // *strings.Reader and returns true. If msg.Body is nil, the length parameter // is unset and the function returns true. Otherwise, false is returned and the // length parameter remains unchanged. func (msg *Message) TryComputeLength() bool { switch v := msg.Body.(type) { case *bytes.Buffer: msg.SetLength(int64(v.Len())) case *bytes.Reader: msg.SetLength(int64(v.Len())) case *strings.Reader: msg.SetLength(int64(v.Len())) case nil: msg.SetLength(0) default: return false } return true } // Intent retrieves the message header intent. func (msg *Message) Intent() string { return msg.Header.Intent } // SetIntent sets the message header intent. func (msg *Message) SetIntent(s string) { msg.Header.Intent = s } // Param retrieves a header parameter. It performs no value validation. func (msg *Message) Param(key string) string { return msg.Header.Parameters[key] } // SetParam sets a header parameter. If the value is empty, the parameter is // unset. It performs no value validation. func (msg *Message) SetParam(key, value string) { if len(value) == 0 { delete(msg.Header.Parameters, key) } else { msg.Header.Parameters[key] = value } } // Length gets the length header parameter (or 0 if it's not set or invalid). func (msg *Message) Length() int64 { n, err := getInt(msg, "length", 0) if err != nil { return 0 } return n } // SetLength sets the length header parameter to n. // // If negative or zero, the parameter is unset. func (msg *Message) SetLength(n int64) { if n <= 0 { msg.SetParam("length", "") } else { setInt(msg, "length", n) } } // Validate validates the message header parameter value format (length). func (msg *Message) Validate() error { _, err := getInt(msg, "length", 0) return err } // Write writes the message header and body to w. func (msg *Message) Write(w io.Writer) error { if err := msg.Header.Write(w); err != nil { return err } if msg.Body != nil { if _, err := io.Copy(w, msg.Body); err != nil { return err } } return nil }