agila_open_tool/tools/v/vlib/net/http/response.v

// Copyright (c) 2019-2024 Alexander Medvednikov. All rights reserved.
// Use of this source code is governed by an MIT license
// that can be found in the LICENSE file.
module http

import compress.gzip
import compress.zlib
import net.http.chunked
import strconv
import strings

// Response represents the result of the request
pub struct Response {
pub mut:
	body         string
	header       Header
	status_code  int
	status_msg   string
	http_version string
}

fn (mut resp Response) free() {
	unsafe { resp.header.free() }
}

// Formats resp to bytes suitable for HTTP response transmission
pub fn (resp Response) bytes() []u8 {
	mut sb := strings.new_builder(resp.response_buffer_cap())
	resp.write_into_builder(mut sb)
	return unsafe { sb.reuse_as_plain_u8_array() }
}

// Formats resp to a string suitable for HTTP response transmission
pub fn (resp Response) bytestr() string {
	mut sb := strings.new_builder(resp.response_buffer_cap())
	resp.write_into_builder(mut sb)
	res := sb.str()
	unsafe { sb.free() }
	return res
}

fn (resp Response) response_buffer_cap() int {
	return resp.body.len + 64 + resp.header.cur_pos * 48
}

fn (resp Response) write_into_builder(mut sb strings.Builder) {
	sb.write_string('HTTP/')
	sb.write_string(resp.http_version)
	sb.write_u8(` `)
	sb.write_decimal(resp.status_code)
	sb.write_u8(` `)
	sb.write_string(resp.status_msg)
	sb.write_string('\r\n')
	resp.header.render_into_sb(mut sb,
		version: resp.version()
	)
	sb.write_string('\r\n')
	sb.write_string(resp.body)
}

// Parse a raw HTTP response into a Response object
pub fn parse_response(resp string) !Response {
	version, status_code, status_msg := parse_status_line(resp.all_before('\r\n'))!
	// Build resp header map and separate the body
	start_idx, end_idx := find_headers_range(resp)!
	header := parse_headers(resp.substr(start_idx, end_idx))!
	mut body := resp.substr(end_idx, resp.len)
	if has_header_token(header.get(.transfer_encoding) or { '' }, 'chunked') {
		body = chunked.decode(body)!
	}
	body = decode_response_body(body, header.get(.content_encoding) or { '' })
	return Response{
		http_version: version
		status_code:  status_code
		status_msg:   status_msg
		header:       header
		body:         body
	}
}

fn has_header_token(header_value string, expected_token string) bool {
	for token in parse_header_tokens(header_value) {
		if token == expected_token.to_lower() {
			return true
		}
	}
	return false
}

fn parse_header_tokens(header_value string) []string {
	mut tokens := []string{}
	for part in header_value.split(',') {
		token := part.all_before(';').trim_space().to_lower()
		if token != '' {
			tokens << token
		}
	}
	return tokens
}

fn decode_response_body(body string, content_encoding string) string {
	if body.len == 0 {
		return body
	}
	encodings := parse_header_tokens(content_encoding)
	if encodings.len == 0 {
		return body
	}
	mut decoded := body.bytes()
	for i := encodings.len - 1; i >= 0; i-- {
		encoding := encodings[i]
		decoded = match encoding {
			'gzip', 'x-gzip' {
				gzip.decompress(decoded) or { return body }
			}
			'deflate' {
				zlib.decompress(decoded) or { return body }
			}
			'identity' {
				decoded
			}
			else {
				return body
			}
		}
	}
	return decoded.bytestr()
}

// parse_status_line parses the first HTTP response line into the HTTP
// version, status code, and reason phrase
fn parse_status_line(line string) !(string, int, string) {
	if line.len < 5 || line[..5].to_lower() != 'http/' {
		return error('response does not start with HTTP/, line: `${line}`')
	}
	data := line.split_nth(' ', 3)
	if data.len != 3 {
		return error('expected at least 3 tokens, but found: ${data.len}')
	}
	version := data[0].substr(5, data[0].len)
	// validate version is 1*DIGIT "." 1*DIGIT
	digits := version.split_nth('.', 3)
	if digits.len != 2 {
		return error('HTTP version malformed, found: `${digits}`')
	}
	for digit in digits {
		strconv.atoi(digit) or {
			return error('HTTP version must contain only integers, found: `${digit}`')
		}
	}
	return version, strconv.atoi(data[1])!, data[2]
}

// cookies parses the Set-Cookie headers into Cookie objects
pub fn (r Response) cookies() []Cookie {
	mut cookies := []Cookie{}
	for cookie in r.header.values(.set_cookie) {
		cookies << parse_cookie(cookie) or { continue }
	}
	return cookies
}

// status parses the status_code and returns a corresponding enum field of Status
pub fn (r Response) status() Status {
	return status_from_int(r.status_code)
}

// set_status sets the status_code and status_msg of the response
pub fn (mut r Response) set_status(s Status) {
	r.status_code = s.int()
	r.status_msg = s.str()
}

// version parses the version
pub fn (r Response) version() Version {
	return match r.http_version {
		'1.0' { .v1_0 }
		'1.1' { .v1_1 }
		'2.0' { .v2_0 }
		else { .unknown }
	}
}

// set_version sets the http_version string of the response
pub fn (mut r Response) set_version(v Version) {
	if v == .unknown {
		r.http_version = ''
		return
	}
	maj, min := v.protos()
	r.http_version = '${maj}.${min}'
}

pub struct ResponseConfig {
pub:
	version Version = .v1_1
	status  Status  = .ok
	header  Header
	body    string
}

// new_response creates a Response object from the configuration. This
// function will add a Content-Length header if body is not empty.
pub fn new_response(conf ResponseConfig) Response {
	mut resp := Response{
		body:   conf.body
		header: conf.header
	}
	if resp.body != '' && !resp.header.contains(.content_length) {
		resp.header.add(.content_length, resp.body.len.str())
	}
	resp.set_status(conf.status)
	resp.set_version(conf.version)
	return resp
}

// find_headers_range returns the start (inclusive) and end (exclusive)
// index of the headers in the string, including the trailing newlines. This
// helper function expects the first line in `data` to be the HTTP status line
// (HTTP/1.1 200 OK).
fn find_headers_range(data string) !(int, int) {
	start_idx := data.index('\n') or { return error('no start index found') } + 1
	mut count := 0
	for i := start_idx; i < data.len; i++ {
		if data[i] == `\n` {
			count++
		} else if data[i] != `\r` {
			count = 0
		}
		if count == 2 {
			return start_idx, i + 1
		}
	}
	return error('no end index found')
}