May 7, 2018

Speeding up your API calls with goroutines

I recently started working on an API client with command line tools for Tenable.io called tenago. Tenable.io is a cloud based vulnerability assessment solution based on Nessus. The application consists of two main components:

  1. A client API

  2. Command line utilities to perform tasks such as querying all your assets, target groups, etc.

Everything was going fine until I wrote a feature that queries all configured scans. The function requires that one GET request is sent to the API server for each configured scan ID. I did this as I wanted to get a list of hostnames configured for each scan. This is what the code looks like for the request:

func queryScans(s *QueryState) *util.ResultTable {
    start := time.Now()
    scansList, err := Client.ListScans()
    if err != nil {
        log.Fatal(err)
    }

    columns := []string{"Scan Name", "Targets"}
    resultTable := util.ResultTable{
        Columns: columns,
    }

    for _, scan := range scansList.Scans {
        if Verbose {
            fmt.Println("[+] Getting details for scan " + scan.Name)
        }

        // Send a GET request to get scan details for each 
        // scan id found. We need to do this in order to get
        // the members for each scan 
        info, err := Client.ScanDetails(string(scan.Id))
        if err != nil {
            log.Fatal(err)
        }

        targetsArray := strings.Split(info.ScanInfo.Targets, ",")
        if s.Hostname != "" {
            if contains(targetsArray, s.Hostname) {
                row := []string{scan.Name, info.ScanInfo.Targets}
                resultTable.Rows = append(resultTable.Rows, row)
            }
        }
    }
    if Verbose{
        util.PrintRuler(Verbose)
        fmt.Printf("%.2fs elapsed\n", time.Since(start).Seconds())
    }
    return &resultTable
}

The code for the GET request (ScanDetails) it’s a basic HTTP request that looks like this:

func (c *Client) ScanDetails(id string) (*ScanDetails, error) {
    path := "/scans/" + id
    req, err := c.newRequest("GET", path, nil)
    if err != nil {
        return nil, err
    }

    result, err := c.do(req)
    if err != nil {
        return nil, err
    }
    var scanDetails ScanDetails
    err = json.Unmarshal(result, &scanDetails)

    if err != nil {
        return nil, err
    }
    return &scanDetails, err
}

Looks simple enough, right? While the code above accomplishes the goal, it is very slow as we have to wait for a response for each to each call to ScanDetails in the loop. It took 36.78 seconds to complete a scan query, as it had to call the ScanDetails function 22 times (one for each scan I had configured):

Slow go function

But we can fix that using goroutines so that each request happens concurrently rather than sequentially. That way we can process the results right away as responses are received from the API server.

func queryScans(s *QueryState) *util.ResultTable{
    start := time.Now()
    
    //Get our list os scans so we can get their IDs
    scansList, err := Client.ListScans()
    if err != nil {
        log.Fatal(err)
    }

    columns := []string{"Scan Name", "Targets"}
    resultTable := util.ResultTable{
        Columns: columns,
    }

    //Create as channel
    ch := make(chan *api.ScanDetails)
    for _, scan := range scansList.Scans {
        if Verbose{
                fmt.Println("[+] Getting details for scan " + scan.Name)
            }
            //Create one goroutine per GET request
            go Client.ScanDetailsWithChannel(string(scan.Id), ch)
    }

    // Process the results of each GET goroutine as they complete
    for range scansList.Scans {
        info := <-ch
        // I added the below to see responses as they are received
        if Verbose{
            fmt.Println("[+] Received data for scan " + info.ScanInfo.Name)
        }
        if s.Hostname != ""{
            if contains(targetsArray, s.Hostname){
                row := []string{info.ScanInfo.Name, info.ScanInfo.Targets}
                resultTable.Rows = append(resultTable.Rows, row)
            }
        }
    }

    if Verbose{
        util.PrintRuler(Verbose)
        fmt.Printf("%.2fs elapsed\n", time.Since(start).Seconds())
    }

    return &resultTable
}

The code for the GET request is only slightly different, as it now expects a channel in the parameters:

func (c *Client) ScanDetailsWithChannel(id string, ch chan <- *ScanDetails) {
    path := "/scans/" + id
    req, err := c.newRequest("GET", path, nil)
    if err != nil {
        log.Fatal(err)
    }

    result, err := c.do(req)
    if err != nil {
        log.Fatal(err)
    }
    var scanDetails ScanDetails
    err = json.Unmarshal(result, &scanDetails)

    if err != nil {
        log.Fatal(err)
    }
    ch <- &scanDetails
}

Nothing too complicated at all, but the speed improvements are noticeable, as now the function does the job in only 6.75 seconds!

Slow go function

Goroutines not only allow you to greatly improve the performance of your code (specially when making multiple API calls in loops) but they are also easy to write. You can learn more about goroutines here

© Dharma of Code 2018