Detecting file format using signature instead of file extension

kubeflow · Mar 7, 2019 · cf7b839 · cf7b839
1 parent ebf9b3b
commit cf7b839
Showing 1 changed file with 8 additions and 11 deletions.
diff --git a/backend/src/apiserver/server/util.go b/backend/src/apiserver/server/util.go
@@ -54,20 +54,20 @@ func loadFile(fileReader io.Reader, maxFileLength int) ([]byte, error) {
 	return pipelineFile[:size], nil
 }
 
-func isSupportedPipelineFormat(fileName string) bool {
-	return isYamlFile(fileName) || isCompressedTarballFile(fileName) || isZipFile(fileName)
+func isSupportedPipelineFormat(fileName string, compressedFile []byte) bool {
+	return isYamlFile(fileName) || isCompressedTarballFile(compressedFile) || isZipFile(compressedFile)
 }
 
 func isYamlFile(fileName string) bool {
 	return strings.HasSuffix(fileName, ".yaml") || strings.HasSuffix(fileName, ".yml")
 }
 
-func isZipFile(fileName string) bool {
-	return strings.HasSuffix(fileName, ".zip")
+func isZipFile(compressedFile []byte) bool {
+	return len(compressedFile) > 2 && compressedFile[0] == '\x50' && compressedFile[1] == '\x4B' //Signature of zip file is "PK"
 }
 
-func isCompressedTarballFile(fileName string) bool {
-	return strings.HasSuffix(fileName, ".tar.gz")
+func isCompressedTarballFile(compressedFile []byte) bool {
+	return len(compressedFile) > 2 && compressedFile[0] == '\x1F' && compressedFile[1] == '\x8B'
 }
 
 func DecompressPipelineTarball(compressedFile []byte) ([]byte, error) {
@@ -113,9 +113,6 @@ func DecompressPipelineZip(compressedFile []byte) ([]byte, error) {
 }
 
 func ReadPipelineFile(fileName string, fileReader io.Reader, maxFileLength int) ([]byte, error) {
-	if !isSupportedPipelineFormat(fileName) {
-		return nil, util.NewInvalidInputError("Unexpected pipeline file format. Support formats are .zip, .tar.gz or YAML.")
-	}
 
 	// Read file into size limited byte array.
 	pipelineFileBytes, err := loadFile(fileReader, maxFileLength)
@@ -127,9 +124,9 @@ func ReadPipelineFile(fileName string, fileReader io.Reader, maxFileLength int)
 	switch {
 	case isYamlFile(fileName):
 		processedFile = pipelineFileBytes
-	case isZipFile(fileName):
+	case isZipFile(pipelineFileBytes):
 		processedFile, err = DecompressPipelineZip(pipelineFileBytes)
-	case isCompressedTarballFile(fileName):
+	case isCompressedTarballFile(pipelineFileBytes):
 		processedFile, err = DecompressPipelineTarball(pipelineFileBytes)
 	default:
 		return nil, util.NewInvalidInputError("Unexpected pipeline file format. Support .zip, .tar.gz or YAML.")