diff --git a/src/chocolatey.resources/helpers/functions/Get-ChocolateyWebFile.ps1 b/src/chocolatey.resources/helpers/functions/Get-ChocolateyWebFile.ps1 index 0b3fb35575..e378807430 100644 --- a/src/chocolatey.resources/helpers/functions/Get-ChocolateyWebFile.ps1 +++ b/src/chocolatey.resources/helpers/functions/Get-ChocolateyWebFile.ps1 @@ -287,7 +287,7 @@ Get-FtpFile if ($url.StartsWith('http:')) { try { $httpsUrl = $url.Replace("http://", "https://") - Get-WebHeaders -Url $httpsUrl -ErrorAction "Stop" | Out-Null + Get-WebHeaders -Url $httpsUrl -Options $options -ErrorAction "Stop" | Out-Null $url = $httpsUrl Write-Warning "Url has SSL/TLS available, switching to HTTPS for download" } @@ -301,7 +301,7 @@ Get-FtpFile $fileFullPath = $fileFullPath -replace '\\chocolatey\\chocolatey\\', '\chocolatey\' $fileDirectory = [System.IO.Path]::GetDirectoryName($fileFullPath) $originalFileName = [System.IO.Path]::GetFileName($fileFullPath) - $fileFullPath = Get-WebFileName -Url $url -DefaultName $originalFileName + $fileFullPath = Get-WebFileName -Url $url -DefaultName $originalFileName -Options $options $fileFullPath = Join-Path $fileDirectory $fileFullPath $fileFullPath = [System.IO.Path]::GetFullPath($fileFullPath) } @@ -324,7 +324,7 @@ Get-FtpFile $headers = @{} if ($url.StartsWith('http')) { try { - $headers = Get-WebHeaders -Url $url -ErrorAction "Stop" + $headers = Get-WebHeaders -Url $url -Options $options -ErrorAction "Stop" } catch { if ($PSVersionTable.PSVersion -lt (New-Object 'Version' 3, 0)) { @@ -333,7 +333,7 @@ Get-FtpFile $originalProtocol = [System.Net.ServicePointManager]::SecurityProtocol [System.Net.ServicePointManager]::SecurityProtocol = [System.Net.SecurityProtocolType]::Ssl3 try { - $headers = Get-WebHeaders -Url $url -ErrorAction "Stop" + $headers = Get-WebHeaders -Url $url -Options $options -ErrorAction "Stop" } catch { Write-Host "Attempt to get headers for $url failed.`n $($_.Exception.Message)" diff --git a/src/chocolatey.resources/helpers/functions/Get-WebFileName.ps1 b/src/chocolatey.resources/helpers/functions/Get-WebFileName.ps1 index 36f5bdf903..4f2488f400 100644 --- a/src/chocolatey.resources/helpers/functions/Get-WebFileName.ps1 +++ b/src/chocolatey.resources/helpers/functions/Get-WebFileName.ps1 @@ -53,6 +53,9 @@ from the url response. The user agent to use as part of the request. Defaults to 'chocolatey command line'. +.PARAMETER Options +OPTIONAL - Specify custom headers. + .PARAMETER IgnoredArguments Allows splatting with arguments that do not apply. Do not use directly. @@ -69,6 +72,7 @@ Get-ChocolateyWebFile [parameter(Mandatory = $false, Position = 0)][string] $url = '', [parameter(Mandatory = $true, Position = 1)][string] $defaultName, [parameter(Mandatory = $false)][string] $userAgent = 'chocolatey command line', + [parameter(Mandatory = $false)][hashtable] $options = @{Headers = @{} }, [parameter(ValueFromRemainingArguments = $true)][Object[]] $ignoredArguments ) @@ -174,6 +178,30 @@ Get-ChocolateyWebFile $request.CookieContainer = New-Object System.Net.CookieContainer $request.UserAgent = $userAgent + if ($options.Headers.Count -gt 0) { + Write-Debug "Setting custom headers" + foreach ($key in $options.Headers.Keys) { + $uri = New-Object -TypeName System.Uri $url + switch ($key) { + 'Accept' { + $request.Accept = $options.Headers[$key] + } + 'Cookie' { + $request.CookieContainer.SetCookies($uri, $options.Headers[$key]) + } + 'Referer' { + $request.Referer = $options.Headers[$key] + } + 'User-Agent' { + $request.UserAgent = $options.Headers[$key] + } + Default { + $request.Headers.Add($key, $options.Headers[$key]) + } + } + } + } + [System.Text.RegularExpressions.Regex]$containsABadCharacter = New-Object Regex("[" + [System.Text.RegularExpressions.Regex]::Escape([System.IO.Path]::GetInvalidFileNameChars() -join '') + "\=\;]"); try { diff --git a/src/chocolatey.resources/helpers/functions/Get-WebHeaders.ps1 b/src/chocolatey.resources/helpers/functions/Get-WebHeaders.ps1 index 8fd96503bf..35ff44422b 100644 --- a/src/chocolatey.resources/helpers/functions/Get-WebHeaders.ps1 +++ b/src/chocolatey.resources/helpers/functions/Get-WebHeaders.ps1 @@ -37,6 +37,9 @@ This is the url to get a request/response from. The user agent to use as part of the request. Defaults to 'chocolatey command line'. +.PARAMETER Options +OPTIONAL - Specify custom headers. + .PARAMETER IgnoredArguments Allows splatting with arguments that do not apply. Do not use directly. @@ -52,6 +55,7 @@ Get-WebFile param( [parameter(Mandatory = $false, Position = 0)][string] $url = '', [parameter(Mandatory = $false, Position = 1)][string] $userAgent = 'chocolatey command line', + [parameter(Mandatory = $false)][hashtable] $options = @{Headers = @{} }, [parameter(ValueFromRemainingArguments = $true)][Object[]] $ignoredArguments ) @@ -133,6 +137,30 @@ Get-WebFile $request.UserAgent = $userAgent } + if ($options.Headers.Count -gt 0) { + Write-Debug "Setting custom headers" + foreach ($key in $options.Headers.Keys) { + $uri = New-Object -TypeName System.Uri $url + switch ($key) { + 'Accept' { + $request.Accept = $options.Headers[$key] + } + 'Cookie' { + $request.CookieContainer.SetCookies($uri, $options.Headers[$key]) + } + 'Referer' { + $request.Referer = $options.Headers[$key] + } + 'User-Agent' { + $request.UserAgent = $options.Headers[$key] + } + Default { + $request.Headers.Add($key, $options.Headers[$key]) + } + } + } + } + Write-Debug "Request Headers:" foreach ($key in $request.Headers) { $value = $request.Headers[$key]; diff --git a/tests/pester-tests/powershell-commands/Get-ChocolateyWebFile.Tests.ps1 b/tests/pester-tests/powershell-commands/Get-ChocolateyWebFile.Tests.ps1 new file mode 100644 index 0000000000..34389e292c --- /dev/null +++ b/tests/pester-tests/powershell-commands/Get-ChocolateyWebFile.Tests.ps1 @@ -0,0 +1,86 @@ +Describe 'Get-ChocolateyWebFile custom header tests' -Tags GetChocolateyWebFile, Cmdlets { + BeforeAll { + Initialize-ChocolateyTestInstall + + $testLocation = Get-ChocolateyTestLocation + Import-Module "$testLocation\helpers\chocolateyInstaller.psm1" + } + + It 'uses custom headers when retrieving web headers and downloading a file' { + $expectedAuthorization = 'Bearer test-token' + $responseBody = 'test response' + $serverPortListener = New-Object System.Net.Sockets.TcpListener ([System.Net.IPAddress]::Loopback, 0) + $serverPortListener.Start() + $serverPort = ([System.Net.IPEndPoint]$serverPortListener.LocalEndpoint).Port + $serverPortListener.Stop() + + $serverJob = Start-Job -ScriptBlock { + param($prefix, $expectedAuthorization, $responseBody, $requestCount) + + $listener = New-Object System.Net.HttpListener + $listener.Prefixes.Add($prefix) + $listener.Start() + Write-Output 'READY' + + $authorizations = @() + for ($requestNumber = 0; $requestNumber -lt $requestCount; $requestNumber++) { + $context = $listener.GetContext() + $authorization = $context.Request.Headers['Authorization'] + $authorizations += $authorization + + $response = $context.Response + if ($authorization -eq $expectedAuthorization) { + $response.StatusCode = 200 + $body = [System.Text.Encoding]::UTF8.GetBytes($responseBody) + $response.ContentLength64 = $body.Length + $response.ContentType = 'application/octet-stream' + $response.Headers.Add('Content-Disposition', 'attachment; filename=downloaded.bin') + $response.OutputStream.Write($body, 0, $body.Length) + } + else { + $response.StatusCode = 401 + } + $response.Close() + } + + $listener.Stop() + $authorizations + } -ArgumentList "http://localhost:$serverPort/", $expectedAuthorization, $responseBody, 4 + + $destination = Join-Path $testLocation 'custom-headers-test.bin' + $downloadedFile = Join-Path $testLocation 'downloaded.bin' + $options = @{ Headers = @{ Authorization = $expectedAuthorization } } + $checksum = ([System.BitConverter]::ToString(([System.Security.Cryptography.SHA256]::Create()).ComputeHash([System.Text.Encoding]::UTF8.GetBytes($responseBody)))).Replace('-', '').ToLowerInvariant() + + try { + $readyDeadline = (Get-Date).AddSeconds(10) + do { + $serverOutput = @(Receive-Job -Job $serverJob -Keep) + if ($serverOutput -contains 'READY') { + break + } + Start-Sleep -Milliseconds 100 + } while ((Get-Date) -lt $readyDeadline) + + $serverOutput | Should -Contain 'READY' + + $headers = Get-WebHeaders -Url "http://localhost:$serverPort/file.bin" -Options $options + $headers['Content-Length'] | Should -Be $responseBody.Length.ToString() + + $result = Get-ChocolateyWebFile -PackageName 'custom-header-test' -FileFullPath $destination -Url "http://localhost:$serverPort/file.bin" -Options $options -Checksum $checksum -ChecksumType sha256 -GetOriginalFileName + $result | Should -Be $downloadedFile + [System.IO.File]::ReadAllText($downloadedFile) | Should -BeExactly $responseBody + + Wait-Job -Job $serverJob -Timeout 30 | Should -Not -BeNullOrEmpty + $authorizations = @(Receive-Job -Job $serverJob | Where-Object { $_ -ne 'READY' }) + $authorizations | Should -HaveCount 4 + $authorizations | Should -Be @($expectedAuthorization, $expectedAuthorization, $expectedAuthorization, $expectedAuthorization) + } + finally { + Stop-Job -Job $serverJob -ErrorAction SilentlyContinue + Remove-Job -Job $serverJob -Force -ErrorAction SilentlyContinue + Remove-Item -Path $destination -Force -ErrorAction SilentlyContinue + Remove-Item -Path $downloadedFile -Force -ErrorAction SilentlyContinue + } + } +}